@alwaysai/device-agent 2.0.0 → 2.0.2-0

package/src/cloud-connection/message-dispatcher.ts

@@ -0,0 +1,33 @@
+import { logger } from '../util/logger';
+
+export interface MessageHandler<T = any> {
+  handle(payload: T, topic?: string): void;
+}
+
+export class MessageDispatcher<T = any> {
+  private handlers: Map<string, MessageHandler<T>> = new Map();
+
+  public registerHandler(topic: string, handler: MessageHandler): void {
+    this.handlers.set(topic, handler);
+  }
+
+  public dispatch(topic: string, payload: T): void {
+    const handler = this.handlers.get(topic);
+    if (handler) {
+      handler.handle(payload, topic);
+    } else {
+      this.handleUnknownMessage(topic, payload);
+    }
+  }
+
+  // Handle unknown or unregistered message topic
+  private handleUnknownMessage(topic: string, payload: T): void {
+    logger.error(
+      `No handler found for topic/type: ${topic} for message ${JSON.stringify(
+        payload,
+        null,
+        2
+      )}`
+    );
+  }
+}

package/src/cloud-connection/publisher.ts

@@ -6,15 +6,16 @@ import {
   getToCloudTopic
 } from '@alwaysai/device-agent-schemas';
 import * as winston from 'winston';
+import { ConnectionManager } from './connection-manager';
 
 export class Publisher {
-  private client: any;
+  private connectionManager: ConnectionManager;
   private clientId: string;
   private readonly toClientTopic: string;
   private readonly toCloudTopic: string;
 
-  constructor(client: any, clientId: string) {
-    this.client = client;
+  constructor(connectionManager: ConnectionManager, clientId: string) {
+    this.connectionManager = connectionManager;
     this.clientId = clientId;
     this.toClientTopic = getToClientTopic(this.clientId);
     this.toCloudTopic = getToCloudTopic(this.clientId);
@@ -34,13 +35,15 @@ export class Publisher {
         2
       )}`
     );
-    this.client.publish(topic, payload, (err: any) => {
-      if (err) {
-        logger.error(
-          `Error publishing message: \nTopic: ${topic}\nMessage: ${payload}\nError: ${err}`
-        );
-      }
-    });
+    this.connectionManager
+      .getIoTDevice()
+      .publish(topic, payload, (err: any) => {
+        if (err) {
+          logger.error(
+            `Error publishing message: \nTopic: ${topic}\nMessage: ${payload}\nError: ${err}`
+          );
+        }
+      });
   }
 
   public publishToCloudWithAck(
@@ -48,19 +51,21 @@ export class Publisher {
     ackNackCallback: CallableFunction
   ) {
     const topic = this.toCloudTopic;
-    this.client.publish(topic, payload, { qos: 1 }, (err: any, resp: any) => {
-      if (err) {
-        logger.error(
-          `Error publishing message: \nTopic: ${topic}\nMessage: ${payload}\nError: ${err}`
-        );
-        ackNackCallback(false);
-      } else if (resp) {
-        logger.debug(
-          `Successfully published message: \nTopic: ${topic}\nMessage: ${payload}`
-        );
-        ackNackCallback(true);
-      }
-    });
+    this.connectionManager
+      .getIoTDevice()
+      .publish(topic, payload, { qos: 1 }, (err: any, resp: any) => {
+        if (err) {
+          logger.error(
+            `Error publishing message: \nTopic: ${topic}\nMessage: ${payload}\nError: ${err}`
+          );
+          ackNackCallback(false);
+        } else if (resp) {
+          logger.debug(
+            `Successfully published message: \nTopic: ${topic}\nMessage: ${payload}`
+          );
+          ackNackCallback(true);
+        }
+      });
   }
 
   public publishDeviceAgentMessage(

package/src/cloud-connection/shadow-handler.test.ts

@@ -5,10 +5,14 @@ import { ShadowHandler } from './shadow-handler';
 import { Logger } from 'winston';
 import { logger } from '../util/logger';
 import { getShadowTopic } from '@alwaysai/device-agent-schemas';
+import { ConnectionManager } from './connection-manager';
 
 jest.mock('../application-control');
 jest.mock('./publisher');
-const mockClient = jest.fn();
+jest.mock('./connection-manager');
+const mockConnectionManager = {
+  getIotDevice: jest.fn()
+} as any as ConnectionManager;
 const clientId = 'test-client';
 const projectId1 = 'test-project';
 const projectId2 = 'test-project-2';
@@ -18,7 +22,7 @@ describe('Test Shadow Handler', () => {
   let shadowHandler: ShadowHandler;
 
   beforeEach(() => {
-    publisher = new Publisher(mockClient, clientId);
+    publisher = new Publisher(mockConnectionManager, clientId);
     shadowHandler = new ShadowHandler(clientId, publisher);
   });
 

package/src/cloud-connection/shadow-handler.ts

@@ -3,13 +3,17 @@ import {
   validateAppConfig
 } from '@alwaysai/app-configuration-schemas';
 import {
+  buildSignedUrlsRequestMessage,
+  buildToClientStatusResponseMessage,
   buildUpdateShadowMessage,
   EnvVars,
   generateTxId,
   getShadowTopic,
   getUpdateDeltaStateFromMessage,
+  keyMirrors,
   ProjectShadowUpdate,
   SecureTunnelShadowUpdate,
+  SignedUrlsRequestPayload,
   validateEnvVarSchemaShadowUpdate,
   validateProjectShadowUpdate
 } from '@alwaysai/device-agent-schemas';
@@ -19,11 +23,19 @@ import {
   ShadowProjectsUpdateAll
 } from '@alwaysai/device-agent-schemas/lib/shadow-schema';
 import { stringifyError } from 'alwaysai/lib/util';
-import { getAllEnvs, readAppCfgFile, setEnv } from '../application-control';
+import {
+  getAllEnvs,
+  readAppCfgFile,
+  setEnv,
+  updateAppCfg
+} from '../application-control';
 import { getSystemInformation } from '../device-control/device-control';
 import { logger } from '../util/logger';
 import { Publisher } from './publisher';
 import { AppConfigModels, getAppCfgModelsDiff } from './shadow';
+import { pruneModels } from '../application-control/models';
+import { MessageHandler } from './message-dispatcher';
+import { BaseHandler } from './base-message-handler';
 
 export type AppConfigUpdate = {
   newAppCfg: AppConfig;
@@ -334,6 +346,10 @@ export class ShadowHandler {
     return [];
   }
 
+  public async initShadows() {
+    await this.updateSystemInfoShadow();
+  }
+
   public async updateSystemInfoShadow() {
     const systemInfo = await getSystemInformation();
     this.publisher.publish(
@@ -440,3 +456,115 @@ export class ShadowHandler {
     );
   }
 }
+
+export class ProjectShadowMessageHandler
+  extends BaseHandler
+  implements MessageHandler<any>
+{
+  public async handle(message: any, topic: string): Promise<void> {
+    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);
+    }
+  }
+
+  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;
+  };
+}

package/src/cloud-connection/transaction-manager.test.ts

@@ -1,18 +1,10 @@
-import {
-  ToClientStatusResponseMessage,
-  generateTxId,
-  keyMirrors
-} from '@alwaysai/device-agent-schemas';
+import { EnvVars, generateTxId } from '@alwaysai/device-agent-schemas';
 import { TransactionManager } from './transaction-manager';
 import { v4 as uuidv4 } from 'uuid';
-import { Publisher } from './publisher';
 import { LiveUpdatesHandler } from './live-updates-handler';
 import { AppConfigUpdate, ShadowUpdate } from './shadow-handler';
-
-const mockClient = {
-  publish: jest.fn()
-};
-const clientId = 'test-client';
+import { AppContent } from './device-agent-message-handler';
+import { AppConfig } from '@alwaysai/app-configuration-schemas';
 
 const mockLiveUpdatesHandler = {
   enable: jest.fn(),
@@ -21,7 +13,6 @@ const mockLiveUpdatesHandler = {
 
 describe('Test Transaction Manager', () => {
   let txnMgr: TransactionManager;
-  let publisher: Publisher;
   const func_complete: () => Promise<boolean> = jest
     .fn()
     .mockResolvedValue(true);
@@ -30,8 +21,7 @@ describe('Test Transaction Manager', () => {
     .mockResolvedValue(false);
 
   beforeEach(() => {
-    publisher = new Publisher(mockClient, clientId);
-    txnMgr = new TransactionManager(publisher, mockLiveUpdatesHandler);
+    txnMgr = new TransactionManager(mockLiveUpdatesHandler);
     jest.clearAllMocks();
   });
 
@@ -58,19 +48,23 @@ describe('Test Transaction Manager', () => {
   test('Start a new transaction which completes in one step', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
+
+    const successFn = jest.fn();
+    const errorFn = jest.fn();
+
     await txnMgr.runTransactionStep({
       func: func_complete,
       projectId,
       txId,
-      start: true
+      start: true,
+      successFn,
+      errorFn
     });
     expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
     expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
     expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
-    const msg: ToClientStatusResponseMessage = JSON.parse(
-      mockClient.publish.mock.calls[0][1]
-    );
-    expect(msg.payload.status).toBe(keyMirrors.statusResponse.success);
+    expect(successFn).toBeCalledTimes(1);
+    expect(errorFn).toBeCalledTimes(0);
   });
 
   test('Start a new transaction and complete in second step', async () => {
@@ -198,6 +192,10 @@ describe('Test Transaction Manager', () => {
   test('Handle error in step function', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
+
+    const successFn = jest.fn();
+    const errorFn = jest.fn();
+
     await txnMgr.runTransactionStep({
       func: jest.fn().mockImplementation(() => {
         throw new Error('Test error!');
@@ -205,15 +203,15 @@ describe('Test Transaction Manager', () => {
       projectId,
       txId,
       start: true,
-      stepName: 'step1'
+      stepName: 'step1',
+      successFn,
+      errorFn
     });
     expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
     expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
     expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
-    const msg: ToClientStatusResponseMessage = JSON.parse(
-      mockClient.publish.mock.calls[0][1]
-    );
-    expect(msg.payload.status).toBe(keyMirrors.statusResponse.failure);
+    expect(successFn).toBeCalledTimes(0);
+    expect(errorFn).toBeCalledTimes(1);
   });
 
   test('Complete ongoing transaction', async () => {
@@ -369,4 +367,37 @@ describe('Test Transaction Manager', () => {
     expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
     expect(txnMgr.getAppCfgUpdateFromTxID(txId)).toEqual(undefined);
   });
+
+  test('store appContent', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    const appCfg: AppConfig = {
+      models: {},
+      scripts: { start: 'python app.py' }
+    };
+    const envVars: EnvVars = { alwaysai: { TEST: '1' } };
+    const cfgUpdate: AppContent = {
+      projectId: projectId,
+      appCfg: appCfg,
+      envVars: envVars
+    };
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
+
+    expect(txnMgr.getAppContentFromTxId(txId)).toEqual(undefined);
+
+    txnMgr.setAppContentToTx(txId, cfgUpdate);
+
+    expect(txnMgr.getProjectFromTransaction(txId)).toEqual(projectId);
+    expect(txnMgr.getTransactionFromProject(projectId)).toEqual(txId);
+    expect(txnMgr.isOngoingTransaction(txId)).toBe(true);
+    expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(true);
+    expect(txnMgr.isAnyOngoingTransaction()).toBe(true);
+    expect(txnMgr.getAppContentFromTxId(txId)).toEqual(cfgUpdate);
+  });
 });

package/src/cloud-connection/transaction-manager.ts

@@ -4,11 +4,11 @@ import {
   keyMirrors
 } from '@alwaysai/device-agent-schemas';
 import { LiveUpdatesHandler } from './live-updates-handler';
-import { Publisher } from './publisher';
 import { logger } from '../util/logger';
 import { keyMirror, stringifyError } from 'alwaysai/lib/util';
 import { CodedError } from '@carnesen/coded-error';
 import { ShadowUpdate } from './shadow-handler';
+import { AppContent } from './device-agent-message-handler';
 
 interface TransactionDetails {
   txId: string;
@@ -18,14 +18,17 @@ interface TransactionDetails {
   update?: string;
   stop?: string;
   appCfgUpdate?: ShadowUpdate;
+  appContent?: AppContent;
 }
 
+export type ErrorFunction = (txid: string, message: string) => void;
+
+export type SuccessFunction = (txId: string) => void;
+
 export class TransactionManager {
   private detailsByTx: Record<string, TransactionDetails> = {};
   private detailsByProject: Record<string, TransactionDetails> = {};
-  private liveUpdatesHandler: LiveUpdatesHandler;
-  private publisher: Publisher;
-
+  private readonly liveUpdatesHandler: LiveUpdatesHandler;
   private async startTransaction(
     txId: string,
     projectId: string,
@@ -91,10 +94,9 @@ export class TransactionManager {
     logger.info(`Updated transaction:\n${JSON.stringify(txDetails, null, 2)}`);
   }
 
-  constructor(publisher: Publisher, liveUpdatesHandler: LiveUpdatesHandler) {
+  constructor(liveUpdatesHandler: LiveUpdatesHandler) {
     this.detailsByTx = {};
     this.detailsByProject = {};
-    this.publisher = publisher;
     this.liveUpdatesHandler = liveUpdatesHandler;
   }
 
@@ -111,9 +113,19 @@ export class TransactionManager {
     start: boolean;
     liveUpdatesPublishFn?: () => Promise<void>;
     stepName?: string;
+    errorFn?: ErrorFunction;
+    successFn?: SuccessFunction;
   }) {
-    const { func, projectId, txId, start, liveUpdatesPublishFn, stepName } =
-      props;
+    const {
+      func,
+      projectId,
+      txId,
+      start,
+      liveUpdatesPublishFn,
+      stepName,
+      errorFn,
+      successFn
+    } = props;
     if (start) {
       await this.startTransaction(
         txId,
@@ -127,31 +139,20 @@ export class TransactionManager {
     try {
       const completed = await func();
       if (completed) {
-        this.completeTransaction(
-          txId,
-          buildToClientStatusResponseMessage(
-            this.publisher.getClientId(),
-            { status: keyMirrors.statusResponse.success },
-            txId
-          )
-        );
+        this.completeTransaction(txId);
+        if (successFn) {
+          successFn(txId);
+        }
       }
     } catch (e) {
       logger.error(
         `Failed to execute cmd for ${projectId}!\n${stringifyError(e)}`
       );
 
-      this.completeTransaction(
-        txId,
-        buildToClientStatusResponseMessage(
-          this.publisher.getClientId(),
-          {
-            status: keyMirrors.statusResponse.failure,
-            message: e.message
-          },
-          txId
-        )
-      );
+      this.completeTransaction(txId);
+      if (errorFn) {
+        errorFn(txId, e.message);
+      }
     }
   }
 
@@ -181,6 +182,10 @@ export class TransactionManager {
     return this.detailsByTx[txId]?.appCfgUpdate;
   }
 
+  public getAppContentFromTxId(txId: string): AppContent | undefined {
+    return this.detailsByTx[txId]?.appContent;
+  }
+
   public setAppCfgUpdateToTx(txId: string, appCfgUpdate: ShadowUpdate) {
     if (this.isOngoingTransaction(txId)) {
       this.detailsByTx[txId].appCfgUpdate = appCfgUpdate;
@@ -190,6 +195,16 @@ export class TransactionManager {
       );
   }
 
+  public setAppContentToTx(txId: string, appContent: AppContent) {
+    if (this.isOngoingTransaction(txId)) {
+      logger.debug(`${txId}: Setting AppContent:${JSON.stringify(appContent)}`);
+      this.detailsByTx[txId].appContent = appContent;
+    } else
+      throw new Error(
+        `Could not set appCfgUpdate, the transaction ${txId} does not exist.`
+      );
+  }
+
   public completeTransaction(
     txId: string,
     messageToPublish?: ToClientStatusResponseMessage
@@ -212,9 +227,5 @@ export class TransactionManager {
       keyMirrors.toClientMessageType.status_response,
       txId
     );
-
-    if (messageToPublish) {
-      this.publisher.publishToClient(messageToPublish);
-    }
   }
 }