@alwaysai/device-agent 2.0.0 → 2.0.2-0

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

@@ -1,878 +1,165 @@
-// eslint-disable-next-line
-const awsIot = require('aws-iot-device-sdk');
 import {
-  AppInstallResponsePayload,
-  AppStateControlPayload,
-  AppVersionControlInstallPayload,
-  AppVersionControlUninstallPayload,
-  DeviceActionPayload,
-  ModelsInstallResponsePayload,
-  SignedUrlsRequestPayload,
-  ToCloudMessage,
-  ToDeviceAgentMessage,
-  buildAppLogsMessage,
-  buildAppStateMessage,
-  buildDeviceStatsMessage,
-  buildSignedUrlsRequestMessage,
   buildToClientStatusResponseMessage,
   getToDeviceTopic,
-  getUpdateDeltaStateFromMessage,
-  keyMirrors,
-  validateSecureTunnelShadowUpdate,
-  validateToDeviceAgentMessage
+  keyMirrors
 } from '@alwaysai/device-agent-schemas';
-import {
-  DEVICE_CERTIFICATE_FILE_PATH,
-  DEVICE_PRIVATE_KEY_FILE_PATH
-} from 'alwaysai/lib/infrastructure';
-import { stringifyError } from 'alwaysai/lib/util';
-import { exec } from 'child_process';
 import { existsSync } from 'fs';
-import { promisify } from 'util';
-import {
-  getAppLogs,
-  installApp,
-  restartApp,
-  startApp,
-  stopApp,
-  uninstallApp,
-  updateAppCfg,
-  updateModelsWithPresignedUrls
-} from '../application-control';
-import { createAppBackup, rollbackApp } from '../application-control/backup';
-import { pruneModels } from '../application-control/models';
-import { reboot } from '../device-control/device-control';
 import { ALWAYSAI_ANALYTICS_PASSTHROUGH } from '../environment';
-import { AgentConfigFile } from '../infrastructure/agent-config';
 import { getBootstrapPrivateKeyFilePath } from '../infrastructure/device-certificate';
 import { migrateFromLegacyCertsAndTokens } from '../infrastructure/legacy-migration/legacy-migration';
 import { requiredConfigFilesPresentAndValid } from '../infrastructure/required-config-checks';
 import { getIoTCoreEndpointUrl } from '../infrastructure/urls';
 import { SecureTunnelHandlerSingleton } from '../secure-tunneling/secure-tunneling';
-import AaiError from '../util/aai-error';
+import { SecureTunnelMessageHandler } from '../secure-tunneling/secure-tunnel-message-handler';
 import { getDeviceAgentVersion } from '../util/check-for-updates';
-import { AWS_ROOT_CERTIFICATE_FILE_PATH } from '../util/directories';
 import { getDeviceUuid } from '../util/get-device-id';
 import { logger } from '../util/logger';
 import sleep from '../util/sleep';
 import { bootstrapProvision } from './bootstrap-provision';
 import { LiveUpdatesHandler } from './live-updates-handler';
-import { getAppStatePayload, getDeviceStatsPayload } from './messages';
 import { PassthroughHandler } from './passthrough-handler';
 import { Publisher } from './publisher';
-import { ShadowHandler, ShadowUpdate } from './shadow-handler';
+import { ShadowHandler, ProjectShadowMessageHandler } from './shadow-handler';
+import { JobHandler } from '../jobs/job-handler';
 import { TransactionManager } from './transaction-manager';
-
-const exec_promise = promisify(exec);
+import { ConnectionManager } from './connection-manager';
+import { DeviceAgentMessageHandler } from './device-agent-message-handler';
+import { HandlerContext } from './base-message-handler';
 
 export class DeviceAgentCloudConnection {
+  private connectionManager: ConnectionManager;
+  private transactionManager: TransactionManager;
+  private liveUpdatesHandler: LiveUpdatesHandler;
   public shadowHandler: ShadowHandler;
   public publisher: Publisher;
-  private liveUpdatesHandler: LiveUpdatesHandler;
-  private txnMgr: TransactionManager;
-  private device = awsIot.device;
 
-  private clientId = getDeviceUuid();
-  private host = getIoTCoreEndpointUrl();
-  private port = 8883;
+  private readonly clientId = getDeviceUuid();
+  private readonly host = getIoTCoreEndpointUrl();
+  private readonly port = 8883;
   private readonly toDeviceTopic = getToDeviceTopic(this.clientId);
-  private readonly secureTunnelNotifyTopic = `$aws/things/${this.clientId}/tunnels/notify`;
   private readonly secureTunnelHandler =
     SecureTunnelHandlerSingleton.getInstance();
 
   constructor() {
-    this.device = awsIot.device({
-      keyPath: DEVICE_PRIVATE_KEY_FILE_PATH,
-      certPath: DEVICE_CERTIFICATE_FILE_PATH,
-      caPath: AWS_ROOT_CERTIFICATE_FILE_PATH,
-      clientId: this.clientId,
-      host: this.host,
-      port: this.port,
-      keepalive: 10 // time before re-connect attempt on dropped connection, default is 400 seconds
-    });
-
-    this.publisher = new Publisher(this.device, this.clientId);
-    this.shadowHandler = new ShadowHandler(this.clientId, this.publisher);
     this.liveUpdatesHandler = new LiveUpdatesHandler();
-    this.txnMgr = new TransactionManager(
-      this.publisher,
-      this.liveUpdatesHandler
+
+    // Initialize & setup the connection
+    this.connectionManager = new ConnectionManager(
+      this.clientId,
+      this.host,
+      this.port
     );
 
-    this.subscribe(this.toDeviceTopic);
-    this.subscribe(this.secureTunnelNotifyTopic);
-    for (const topic of this.shadowHandler.projectShadowTopics) {
-      this.subscribe(topic);
-    }
-    this.subscribe(this.shadowHandler.shadowTopics.secureTunnel.updateDelta);
-    this.subscribe(this.shadowHandler.shadowTopics.secureTunnel.deleteAccepted);
+    this.publisher = new Publisher(this.connectionManager, this.clientId);
+    this.shadowHandler = new ShadowHandler(this.clientId, this.publisher);
 
-    this.setupHandlers();
-  }
+    this.connectionManager.connect(() => {
+      void this.shadowHandler.initShadows();
+      this.publisher.publish(JOB_HANDLER_TOPICS.START_NEXT, JSON.stringify({}));
+    });
 
-  /*=================================================================
-                            Public interface
-  =================================================================*/
+    this.transactionManager = new TransactionManager(this.liveUpdatesHandler);
 
-  public getClientId(): string {
-    return this.clientId;
-  }
-
-  public isCmdInProgress(projectId: string): boolean {
-    return this.txnMgr.isOngoingTransactionForProjectID(projectId);
-  }
+    // Construct a HandlerContext used by all the message handlers
+    const handlerContext: HandlerContext = {
+      clientId: this.clientId,
+      txnMgr: this.transactionManager,
+      publisher: this.publisher,
+      shadowHandler: this.shadowHandler,
+      liveUpdatesHandler: this.liveUpdatesHandler,
+      secureTunnelHandler: this.secureTunnelHandler
+    };
 
-  public async handleMessage(topic: string, message: any) {
-    logger.debug(
-      `Received message: ${JSON.stringify({ topic, message }, null, 2)}`
+    // Instantiate & register message handlers for Project Shadow topics
+    const projectShadowMessageHandler = new ProjectShadowMessageHandler(
+      handlerContext
     );
-    // ProjectShadow messages
-    if (this.shadowHandler.projectShadowTopics.includes(topic)) {
-      await this.handleProjectShadowMessage(topic, message);
-    } else if (topic === this.toDeviceTopic) {
-      await this.handleDeviceAgentMessage({
+    this.shadowHandler.projectShadowTopics.forEach((topic) => {
+      this.connectionManager.registerHandler(
         topic,
-        message
-      });
-      // SecureTunnelNotify messages
-    } else if (topic === this.secureTunnelNotifyTopic) {
-      await this.secureTunnelHandler.secureTunnelNotifyHandler(message);
-      // SecureTunnel messages
-    } else if (
-      topic === this.shadowHandler.shadowTopics.secureTunnel.updateDelta
-    ) {
-      await this.handleSecureTunnelMessage(message);
-    } else if (
-      topic === this.shadowHandler.shadowTopics.secureTunnel.deleteAccepted
-    ) {
-      logger.info(`Received secure tunnel deleteAccepted: ${message}`);
-      await this.secureTunnelHandler.destroy();
-    } else {
-      logger.error(`Unexpected topic, ignoring! ${topic}`);
-    }
-  }
-
-  public async stop() {
-    // 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();
-  }
-
-  /*=================================================================
-                            Private interface
-  =================================================================*/
-
-  private setupHandlers() {
-    this.device.on('connect', (connack: any) => {
-      logger.info('Device Agent has connected to the cloud');
-      // FIXME: EI-709 Skip this request for now to prevent kicking off another
-      // shadow update process if IoT Core disconnect occurs during app config update
-      //this.shadowHandler.getShadowUpdates();
-      void this.shadowHandler.updateSystemInfoShadow();
-    });
-
-    this.device.on('disconnect', () => {
-      logger.warn('Device Agent has been disconnected from the cloud');
-    });
-
-    this.device.on('reconnect', () => {
-      logger.info(
-        `Device Agent attempting to re-connect ${new Date().toLocaleString()}`
+        projectShadowMessageHandler
       );
     });
 
-    this.device.on('error', function (e) {
-      logger.error(`Error connecting to cloud!\n${stringifyError(e)}`);
-    });
-
-    this.device.on('message', async (topic: string, payload: string) => {
-      try {
-        const jsonPacket = JSON.parse(payload);
-        await this.handleMessage(topic, jsonPacket);
-      } catch (e) {
-        logger.error(`Error parsing message!\n${stringifyError(e)}`);
-      }
-    });
-
-    this.device.on('offline', () => {
-      logger.warn(`Device Agent is offline ${new Date().toLocaleString()}`);
-      void this.logConnectionInfo();
-    });
-  }
-
-  private async logConnectionInfo() {
-    try {
-      /**
-       * We're using the 'netcat' or 'nc' command to test the connection to the IoT Core endpoint.
-       * This command doesn't always exit (see below), so
-       * we use timeout to break out of the prompt
-       * and catch the resulting error/parse the resulting stderr
-       *
-       * Sample command for current host and port:
-       * nc -zv -w 1 a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com 8883
-       *
-       * Sample output when port is not blocked and host is reachable:
-       * $ nc  -zv -w 1 a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com 443
-       * Connection to a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com 443 port [tcp/https] succeeded!
-       *
-       *
-       * Sample output when port is blocked (will repeatedly try until ctrl-C out):
-       * $ nc -zv -w 1 a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com 8883
-       * nc: connect to a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com port 8883 (tcp) timed out: Operation now in progress
-       * nc: connect to a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com port 8883 (tcp) timed out: Operation now in progress
-       * nc: connect to a3tzi5g7sq5zsj-ats.iot.us-west-2.amazonaws.com port 8883 (tcp) timed out: Operation now in progress
-       * ^C
-       *
-       *
-       * Sample command/output when the port isn't enable on that host:
-       * $ nc -zv -w 1 localhost 8883
-       * nc: connect to localhost port 8883 (tcp) failed: Connection refused
-       */
-      await exec_promise(`nc -zv -w 1 ${this.host} ${this.port}`, {
-        timeout: 2000
-      });
-    } catch (err) {
-      const output = JSON.stringify(err['stderr']);
-      if (output.indexOf('not known') !== -1) {
-        logger.warn(
-          'Iot Core endpoint appears to be unreachable, internet connection may be unstable or the host may be down.'
-        );
-      } else if (output.indexOf('timed out') !== -1) {
-        logger.warn(
-          `Internet connection appears fine, however the endpoint was not reachable on the current connection port: ${this.port}\nPlease check if a firewall is in place.`
-        );
-      } else if (output.indexOf('refused') !== -1) {
-        logger.warn(
-          `The connection was refused, likely ${this.host} is not running a service on ${this.port}.`
-        );
-      } else {
-        logger.warn(
-          `Output from checking connection to ${this.host} on ${this.port}: ${output}`
-        );
-      }
-    }
-  }
-
-  private async handleDeviceAgentMessage({
-    topic,
-    message
-  }: {
-    topic: string;
-    message: ToDeviceAgentMessage;
-  }) {
-    const valid = validateToDeviceAgentMessage(message);
-    if (!valid) {
-      logger.error(
-        `Error validating message: ${JSON.stringify(
-          { topic, message, errors: validateToDeviceAgentMessage.errors },
-          null,
-          2
-        )}`
-      );
-      // TODO: Send generic error response
-      return;
-    }
-    const txId = message.txId;
-    const {
-      app_state_control,
-      app_version_control,
-      live_state_updates,
-      app_install_response,
-      models_install_response,
-      status_response,
-      device_action
-    } = keyMirrors.toDeviceAgentMessageType;
-    switch (message.messageType) {
-      case app_state_control: {
-        // txId sent from cloud, just need to continue it
-        const payload = message.payload;
-        const projectId = payload.projectId;
-
-        try {
-          await this.txnMgr.runTransactionStep({
-            func: () => this.handleAppStateControl(message.payload),
-            projectId,
-            txId,
-            start: true,
-            liveUpdatesPublishFn: async () =>
-              this.publisher.publishToClient(
-                buildToClientStatusResponseMessage(
-                  this.clientId,
-                  { status: keyMirrors.statusResponse.in_progress },
-                  txId
-                ),
-                logger.silly
-              ),
-            stepName: payload.baseCommand
-          });
-        } catch (e) {
-          logger.error(
-            `Error processing application state control request for ${projectId}!\n${stringifyError(
-              e
-            )}`
-          );
-        }
-
-        break;
-      }
-      case app_version_control: {
-        // txId sent from cloud, just need to continue it
-        const payload = message.payload;
-        const projectId = payload.projectId;
-        try {
-          await this.txnMgr.runTransactionStep({
-            func: () => this.handleAppVersionControl(payload, txId),
-            projectId,
-            txId,
-            start: true,
-            liveUpdatesPublishFn: async () =>
-              this.publisher.publishToClient(
-                buildToClientStatusResponseMessage(
-                  this.clientId,
-                  { status: keyMirrors.statusResponse.in_progress },
-                  txId
-                ),
-                logger.silly
-              ),
-            stepName: payload.baseCommand
-          });
-        } catch (e) {
-          logger.error(
-            `Error processing application install request for ${projectId}!\n${stringifyError(
-              e
-            )}`
-          );
-        }
-
-        break;
-      }
-      case live_state_updates: {
-        const { deviceStats, appState, appLogs } = message.payload;
-
-        if (deviceStats !== undefined) {
-          if (deviceStats) {
-            await this.liveUpdatesHandler.enable(
-              keyMirrors.toClientMessageType.device_stats,
-              async () =>
-                this.publisher.publishToClient(
-                  buildDeviceStatsMessage(
-                    this.clientId,
-                    await getDeviceStatsPayload(),
-                    txId
-                  ),
-                  logger.silly
-                )
-            );
-          } else {
-            this.liveUpdatesHandler.disable(
-              keyMirrors.toClientMessageType.device_stats
-            );
-          }
-        }
-
-        if (appState !== undefined) {
-          if (appState) {
-            await this.liveUpdatesHandler.enable(
-              keyMirrors.toClientMessageType.app_state,
-              async () =>
-                this.publisher.publishToClient(
-                  buildAppStateMessage(
-                    this.clientId,
-                    await getAppStatePayload(),
-                    txId
-                  ),
-                  logger.silly
-                )
-            );
-          } else {
-            this.liveUpdatesHandler.disable(
-              keyMirrors.toClientMessageType.app_state
-            );
-          }
-        }
-
-        if (appLogs !== undefined) {
-          if (appLogs.toggle) {
-            await this.liveUpdatesHandler.startStream(
-              appLogs.projectId,
-              async () =>
-                await getAppLogs({
-                  projectId: appLogs.projectId,
-                  args: ['--tail', '100', '--no-log-prefix']
-                }),
-              async (logChunk: string) =>
-                this.publisher.publishToClient(
-                  buildAppLogsMessage(
-                    this.clientId,
-                    {
-                      projectId: appLogs.projectId,
-                      logChunk
-                    },
-                    txId
-                  ),
-                  logger.silly
-                )
-            );
-          } else {
-            this.liveUpdatesHandler.stopStream(appLogs.projectId);
-          }
-        }
-        break;
-      }
-      case app_install_response: {
-        const payload = message.payload;
-        const { projectId } = payload.appInstallResponse;
-        if (txId !== this.txnMgr.getTransactionFromProject(projectId)) {
-          throw new Error(
-            `App install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnMgr.getTransactionFromProject(
-              projectId
-            )})!`
+    // Instantiate & register message handlers for to-device and secureTunnel topics
+    this.connectionManager.registerHandler(
+      this.toDeviceTopic,
+      new DeviceAgentMessageHandler(
+        handlerContext,
+        (txId: string, errorMsg: string) => {
+          const msg = buildToClientStatusResponseMessage(
+            this.publisher.getClientId(),
+            {
+              status: keyMirrors.statusResponse.failure,
+              message: errorMsg
+            },
+            txId
           );
-        }
-        await this.txnMgr.runTransactionStep({
-          func: () => this.handleAppInstallCloudResponsePayload(payload),
-          projectId,
-          txId,
-          start: false,
-          stepName: message.messageType
-        });
-        break;
-      }
-      case 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 payload = message.payload;
-        const { projectId } = payload.modelsInstallResponse;
-        if (txId !== this.txnMgr.getTransactionFromProject(projectId)) {
-          throw new Error(
-            `Model install response received a message for a transaction ID ${txId} that is not currently underway (${this.txnMgr.getTransactionFromProject(
-              projectId
-            )})!`
+          this.publisher.publishToClient(msg);
+        },
+        (txId: string) => {
+          const msg = buildToClientStatusResponseMessage(
+            this.publisher.getClientId(),
+            { status: keyMirrors.statusResponse.success },
+            txId
           );
+          this.publisher.publishToClient(msg);
         }
-        await this.txnMgr.runTransactionStep({
-          func: () =>
-            this.handleModelsInstallCloudResponsePayload(payload, txId),
-          projectId,
-          txId,
-          start: false,
-          stepName: message.messageType
-        });
-        break;
-      }
-      case status_response: {
-        const { failure } = keyMirrors.statusResponse;
-        if (message.payload.status === failure) {
-          this.txnMgr.completeTransaction(
-            txId,
-            buildToClientStatusResponseMessage(
-              this.clientId,
-              {
-                status: keyMirrors.statusResponse.failure,
-                message: message.payload.message
-              },
-              txId
-            )
-          );
-        }
-        break;
-      }
-      case device_action: {
-        try {
-          this.publisher.publishToClient(
-            buildToClientStatusResponseMessage(
-              this.clientId,
-              {
-                status: keyMirrors.statusResponse.in_progress
-              },
-              txId
-            )
-          );
-
-          await this.handleDeviceAction(message.payload);
-
-          this.publisher.publishToClient(
-            buildToClientStatusResponseMessage(
-              this.clientId,
-              {
-                status: keyMirrors.statusResponse.success
-              },
-              txId
-            )
-          );
-        } catch (e) {
-          logger.error(
-            `There was a problem performing device action '${
-              message.payload.action
-            }'!\n${stringifyError(e)}`
-          );
-          this.publisher.publishToClient(
-            buildToClientStatusResponseMessage(
-              this.clientId,
-              {
-                status: keyMirrors.statusResponse.failure,
-                message: e.message
-              },
-              txId
-            )
-          );
-        }
-        break;
-      }
-      default:
-        logger.error(
-          `Invalid client message: '${JSON.stringify(
-            { topic, message, txId },
-            null,
-            2
-          )}'`
-        );
-    }
-  }
-
-  private handleAppStateControl = async (
-    payload: AppStateControlPayload
-  ): Promise<boolean> => {
-    const { baseCommand, projectId } = payload;
-    switch (baseCommand) {
-      case keyMirrors.appStateControl.start:
-        await startApp({ projectId });
-        break;
-      case keyMirrors.appStateControl.stop:
-        await stopApp({ projectId });
-        break;
-      case keyMirrors.appStateControl.restart:
-        await restartApp({ projectId });
-        break;
-    }
-    return true;
-  };
-
-  private handleAppVersionControl = async (
-    payload:
-      | AppVersionControlInstallPayload
-      | AppVersionControlUninstallPayload,
-    txId: string
-  ): Promise<boolean> => {
-    switch (payload.baseCommand) {
-      case keyMirrors.appVersionControl.install: {
-        const { projectId, appReleaseHash } = payload;
-
-        const signedUrlsRequestPayload: SignedUrlsRequestPayload = {
-          signedUrlsRequest: {
-            projectId,
-            appReleaseHash
-          }
-        };
-        const message = buildSignedUrlsRequestMessage(
-          this.clientId,
-          signedUrlsRequestPayload,
-          txId
-        );
-        await this.publishCloudRequest(message);
-        return false;
-      }
-      case keyMirrors.appVersionControl.uninstall: {
-        const { projectId } = payload;
-        await this.atomicApplicationUninstall(projectId);
-        return true;
-      }
-      default:
-        logger.warn(
-          `Ignore App Version Control packet: ${JSON.stringify(
-            payload,
-            null,
-            2
-          )}`
-        );
-        return true;
-    }
-  };
-
-  private handleAppInstallCloudResponsePayload = async (
-    payload: AppInstallResponsePayload
-  ): Promise<boolean> => {
-    const {
-      projectId,
-      appReleaseHash,
-      appInstallPayload,
-      modelsInstallPayload
-    } = payload.appInstallResponse;
-    const signedUrlsPayload = {
-      appInstallPayload,
-      modelsInstallPayload
-    };
-    await this.atomicApplicationUpdate(async () => {
-      this.shadowHandler.clearProjectShadow(projectId);
-      await installApp({ projectId, appReleaseHash, signedUrlsPayload });
-    }, projectId);
-    return true;
-  };
-
-  private handleModelsInstallCloudResponsePayload = async (
-    payload: ModelsInstallResponsePayload,
-    txId: string
-  ): Promise<boolean> => {
-    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;
-    if (appCfgUpdate) {
-      await this.atomicApplicationUpdate(
-        async () =>
-          await updateModelsWithPresignedUrls({
-            projectId,
-            modelInstallPayloads: payload.modelsInstallResponse.newModels,
-            newAppCfg: appCfgUpdate.newAppCfg
-          }),
-        projectId
-      );
-    }
+    const secureTunnelMessageHandler = new SecureTunnelMessageHandler(
+      handlerContext
+    );
+    this.connectionManager.registerHandler(
+      secureTunnelMessageHandler.getNotifyTopic(),
+      secureTunnelMessageHandler
+    );
 
-    if (envVarUpdate) {
-      await this.atomicApplicationUpdate(
-        async () =>
-          await this.shadowHandler.updateProjectEnvVars({
-            projectId,
-            envVars: envVarUpdate.envVars
-          }),
-        projectId,
-        true
-      );
-    }
+    this.connectionManager.registerHandler(
+      this.shadowHandler.shadowTopics.secureTunnel.updateDelta,
+      secureTunnelMessageHandler
+    );
+    this.connectionManager.registerHandler(
+      this.shadowHandler.shadowTopics.secureTunnel.deleteAccepted,
+      secureTunnelMessageHandler
+    );
 
-    return true;
-  };
+    const jobHandler = new JobHandler(handlerContext);
 
-  private async handleDeviceAction(payload: DeviceActionPayload) {
-    const { system_restart } = keyMirrors.deviceAction;
-    switch (payload.action) {
-      case system_restart: {
-        await reboot();
-        break;
-      }
-      default: {
-        logger.info(
-          `Unrecognized device action requested: '${payload.action}'.`
-        );
-      }
-    }
-  }
+    const JOB_HANDLER_TOPICS = jobHandler.getJobTopic();
 
-  private async publishCloudRequest(message: ToCloudMessage) {
-    this.publisher.publishToCloud(message);
-  }
+    this.connectionManager.registerHandler(
+      JOB_HANDLER_TOPICS.NOTIFY_NEXT,
+      jobHandler
+    );
 
-  private subscribe(topic: string) {
-    logger.info(`Subscribing to ${topic}`);
-    this.device.subscribe(topic);
+    this.connectionManager.registerHandler(
+      JOB_HANDLER_TOPICS.START_NEXT_ACCEPTED,
+      jobHandler
+    );
   }
 
-  private async atomicApplicationUninstall(projectId: string) {
-    try {
-      await uninstallApp({ projectId });
-      this.shadowHandler.clearProjectShadow(projectId);
-    } catch (e) {
-      logger.error(`Failed to uninstall ${projectId}!\n${stringifyError(e)}`);
-      throw e;
-    }
+  /*=================================================================
+                            Public interface
+  =================================================================*/
+  public getClientId(): string {
+    return this.clientId;
   }
 
-  // eslint-disable-next-line
-  private async atomicApplicationUpdate<F extends () => any>(
-    func: F,
-    projectId: string,
-    skipUpdateShadow?: boolean
-  ): Promise<ReturnType<F>> {
-    if (await AgentConfigFile().isAppPresent({ projectId })) {
-      // Reject the application update if app is present but not ready
-      if (!(await AgentConfigFile().isAppReady({ projectId }))) {
-        throw new Error('Application already has installation in progress!');
-      }
-
-      // Try to create a backup, so that there is one available if something goes wrong in the next try:catch.
-      try {
-        await createAppBackup({ projectId });
-      } catch (e) {
-        logger.error(
-          `Could not create a backup for the project: ${projectId}!\n${stringifyError(
-            e
-          )}`
-        );
-      }
-    }
-
-    try {
-      const out: ReturnType<F> = await func();
-      if (!skipUpdateShadow)
-        await this.shadowHandler.updateProjectShadow(projectId);
-      return out;
-    } catch (errorAppUpdate) {
-      logger.error(
-        `Failed to update ${projectId}!\n${stringifyError(errorAppUpdate)}`
-      );
-      // If something goes wrong, first try to rollback
-      try {
-        await rollbackApp({ projectId });
-      } catch (errorRollbackApp) {
-        logger.error(
-          `Application rollback failed for ${projectId}!\n${stringifyError(
-            errorRollbackApp
-          )}`
-        );
-        // and if that fails, uninstall the app as a last resort.
-        try {
-          await this.atomicApplicationUninstall(projectId);
-        } catch {
-          // atomicApplicationUninstall logs failure, so there's nothing to do here.
-        }
-        throw new AaiError(
-          'Application update and rollback failed, uninstalled the application!',
-          { cause: errorAppUpdate }
-        );
-      }
-      throw new Error(
-        'Application update failed, rolled the application back!',
-        { cause: errorAppUpdate }
-      );
-    }
+  public isCmdInProgress(projectId: string): boolean {
+    return this.transactionManager.isOngoingTransactionForProjectID(projectId);
   }
 
-  private handleProjectShadowConfigUpdate = async (
-    update: ShadowUpdate,
-    txId: string
-  ): Promise<boolean> => {
-    const { projectId, appCfgUpdate, envVarUpdate } = update;
-
-    if (
-      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(
-          updatedModels
-        )}`
-      );
-      const modelsOnlyUrlsRequestPayload: SignedUrlsRequestPayload = {
-        modelsOnlyUrlsRequest: {
-          projectId,
-          models: updatedModels
-        }
-      };
-      const message = buildSignedUrlsRequestMessage(
-        this.clientId,
-        modelsOnlyUrlsRequestPayload,
-        txId
-      );
-      this.publisher.publishToCloud(message);
-
-      this.txnMgr.setAppCfgUpdateToTx(txId, update);
-
-      return false;
-    }
-
-    if (appCfgUpdate) {
-      await this.atomicApplicationUpdate(async () => {
-        await pruneModels({
-          projectId,
-          appCfg: appCfgUpdate.newAppCfg
-        });
-        await updateAppCfg({
-          projectId,
-          newAppCfg: appCfgUpdate.newAppCfg
-        });
-      }, projectId);
-    }
-
-    if (envVarUpdate) {
-      await this.atomicApplicationUpdate(
-        async () =>
-          await this.shadowHandler.updateProjectEnvVars({
-            projectId,
-            envVars: envVarUpdate.envVars
-          }),
-        projectId,
-        true
-      );
-    }
-    return true;
-  };
-
-  private async handleProjectShadowMessage(topic: string, message: any) {
-    const shadowUpdates = await this.shadowHandler.handleProjectShadow({
-      topic,
-      payload: message,
-      clientToken: message.clientToken
-    });
-    if (shadowUpdates.length) {
-      const shadowUpdatePromises: Promise<void>[] = [];
-      for (const shadowUpdate of shadowUpdates) {
-        const projectId = shadowUpdate.projectId;
-        const txId = shadowUpdate.txId;
-        shadowUpdatePromises.push(
-          this.txnMgr
-            .runTransactionStep({
-              func: () =>
-                this.handleProjectShadowConfigUpdate(shadowUpdate, txId),
-              projectId,
-              txId,
-              start: true,
-              liveUpdatesPublishFn: async () =>
-                this.publisher.publishToClient(
-                  buildToClientStatusResponseMessage(
-                    this.clientId,
-                    { status: keyMirrors.statusResponse.in_progress },
-                    txId
-                  ),
-                  logger.silly
-                ),
-              stepName: topic
-            })
-            .catch((e) => {
-              logger.error(
-                `There was an issue updating project shadow config for ${projectId}!\n${stringifyError(
-                  e
-                )}`
-              );
-            })
-        );
-      }
-
-      await Promise.all(shadowUpdatePromises);
-    }
+  public async stop() {
+    // 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.connectionManager.disconnect();
   }
-
-  public async handleSecureTunnelMessage(payload: any): Promise<void> {
-    logger.info(`Received secure tunnel update: ${JSON.stringify(payload)}`);
-    const state = getUpdateDeltaStateFromMessage(payload);
-    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);
+  public async handleMessage(topic: string, message: any) {
+    this.connectionManager.dispatch(topic, message);
   }
 }