@alwaysai/device-agent 1.4.0 → 2.0.0

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

@@ -3,27 +3,27 @@ import {
   validateAppConfig
 } from '@alwaysai/app-configuration-schemas';
 import {
+  buildUpdateShadowMessage,
   EnvVars,
-  getAllEnvs,
-  readAppCfgFile,
-  setEnv
-} from '../application-control';
+  generateTxId,
+  getShadowTopic,
+  getUpdateDeltaStateFromMessage,
+  ProjectShadowUpdate,
+  SecureTunnelShadowUpdate,
+  validateEnvVarSchemaShadowUpdate,
+  validateProjectShadowUpdate
+} from '@alwaysai/device-agent-schemas';
+import {
+  buildBaseShadowMessage,
+  DeviceAgentStatusShadowUpdate,
+  ShadowProjectsUpdateAll
+} from '@alwaysai/device-agent-schemas/lib/shadow-schema';
+import { stringifyError } from 'alwaysai/lib/util';
+import { getAllEnvs, readAppCfgFile, setEnv } 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 {
-  generateTxId,
-  validateProjectShadowUpdate,
-  buildBaseShadowMessage,
-  buildUpdateProjectShadowMessage,
-  buildUpdateSystemInfoShadowMessage,
-  getShadowTopic,
-  ShadowProjectsUpdateAll,
-  getDesiredFromMessage,
-  ShadowTopics,
-  ProjectShadowUpdate
-} from '@alwaysai/device-agent-schemas';
 
 export type AppConfigUpdate = {
   newAppCfg: AppConfig;
@@ -107,15 +107,15 @@ export class ShadowHandler {
     txId: string;
     projectId: string;
   }): Promise<AppConfigUpdate | null> {
-    let appCfgUpdate: any;
     let newAppCfg: any;
-
     // Handle errors and validation
     try {
       newAppCfg = JSON.parse(appConfig);
-    } catch (error) {
+    } catch (e) {
       logger.error(
-        `Could not parse the appConfig for transaction ${txId}!\n${error}`
+        `Could not parse the appConfig for transaction ${txId}!\n${stringifyError(
+          e
+        )}`
       );
       return null;
     }
@@ -139,14 +139,47 @@ export class ShadowHandler {
       projectId
     });
 
+    const appCfgUpdate: AppConfigUpdate = { newAppCfg };
     if (updatedModels && Object.keys(updatedModels).length) {
-      appCfgUpdate = { newAppCfg, updatedModels };
-    } else {
-      appCfgUpdate = { newAppCfg };
+      appCfgUpdate.updatedModels = updatedModels;
     }
     return appCfgUpdate;
   }
 
+  private async generateEnvVarsUpdate({
+    envVars,
+    txId,
+    projectId
+  }: {
+    envVars: string;
+    txId: string;
+    projectId: string;
+  }): Promise<EnvVarUpdate | null> {
+    let newEnvVars: any;
+    try {
+      newEnvVars = JSON.parse(envVars);
+    } catch (e) {
+      logger.error(
+        `Could not parse the environment variables for transaction ${txId}!\n${stringifyError(
+          e
+        )}`
+      );
+      return null;
+    }
+
+    if (!validateEnvVarSchemaShadowUpdate(newEnvVars)) {
+      logger.error(
+        `Received invalid environment variables update for ${projectId}!\n${JSON.stringify(
+          validateEnvVarSchemaShadowUpdate.errors,
+          null,
+          2
+        )}`
+      );
+      return null;
+    }
+    return { envVars: newEnvVars };
+  }
+
   private async processProjectShadowUpdates({
     delta
   }: {
@@ -220,7 +253,14 @@ export class ShadowHandler {
         `Found a delta for app environment variable shadow. Updating ${projectId}`
       );
       const envVars = projectDelta.envVars;
-      shadowUpdate.envVarUpdate = { envVars };
+      const envVarsUpdate = await this.generateEnvVarsUpdate({
+        envVars,
+        txId,
+        projectId
+      });
+      if (envVarsUpdate) {
+        shadowUpdate.envVarUpdate = envVarsUpdate;
+      }
     }
 
     return shadowUpdate.appCfgUpdate || shadowUpdate.envVarUpdate
@@ -244,24 +284,14 @@ export class ShadowHandler {
       throw Error(`Topic ${topic} is not in the ${this.projectShadowTopics}`);
     }
     switch (topic) {
-      case this.shadowTopics.project.updateAccepted: {
-        if (clientToken === this.clientId) {
+      case this.shadowTopics.project.updateDelta: {
+        const delta = getUpdateDeltaStateFromMessage(payload);
+        if (!delta) {
           logger.debug(
-            `Ignoring message as it was caused by Device Agent itself: ${JSON.stringify(
-              { topic, payload },
-              null,
-              2
-            )}`
-          );
-          break;
-        }
-        const desired = getDesiredFromMessage(payload);
-        if (!desired) {
-          logger.debug(
-            `No desired state found in message: ${JSON.stringify(payload)}`
+            `No delta state found in message: ${JSON.stringify(payload)}`
           );
         } else {
-          return await this.processProjectShadowUpdates({ delta: desired });
+          return await this.processProjectShadowUpdates({ delta });
         }
         break;
       }
@@ -276,21 +306,18 @@ export class ShadowHandler {
           );
           break;
         }
-        if (payload.state.delta) {
-          return await this.processProjectShadowUpdates({
-            delta: payload.state.delta
-          });
-        } else {
-          logger.info(
-            `No delta in projects.getAccepted in named shadow '${
-              topic.split('/')[5]
-            }'`
+        const delta = getUpdateDeltaStateFromMessage(payload);
+        if (!delta) {
+          logger.debug(
+            `No delta state found in message: ${JSON.stringify(payload)}`
           );
+        } else {
+          return await this.processProjectShadowUpdates({ delta });
         }
         break;
       }
       case this.shadowTopics.project.getRejected:
-      case this.shadowTopics.project.updateDelta:
+      case this.shadowTopics.project.updateAccepted:
       case this.shadowTopics.project.updateRejected: {
         // Not handling these for now
         break;
@@ -312,7 +339,10 @@ export class ShadowHandler {
     this.publisher.publish(
       getShadowTopic(this.clientId, 'system-info', 'update'),
       JSON.stringify(
-        buildUpdateSystemInfoShadowMessage(this.clientId, systemInfo)
+        buildUpdateShadowMessage({
+          clientId: this.clientId,
+          reported: systemInfo
+        })
       )
     );
   }
@@ -324,13 +354,13 @@ export class ShadowHandler {
     const toReport: ShadowProjectsUpdateAll = {
       [projectId]: {
         appConfig: JSON.stringify(appCfg),
-        envVars
+        envVars: JSON.stringify(envVars)
       }
     };
     this.publisher.publish(
       getShadowTopic(this.clientId, 'projects', 'update'),
       JSON.stringify(
-        buildUpdateProjectShadowMessage({
+        buildUpdateShadowMessage({
           clientId: this.clientId,
           reported: toReport
         })
@@ -346,15 +376,15 @@ export class ShadowHandler {
     envVars: EnvVars;
   }) {
     await setEnv({ projectId, envVars });
-
+    const currentEnvs = await getAllEnvs({ projectId });
     this.publisher.publish(
       getShadowTopic(this.clientId, 'projects', 'update'),
       JSON.stringify(
-        buildUpdateProjectShadowMessage({
+        buildUpdateShadowMessage({
           clientId: this.clientId,
           reported: {
             [projectId]: {
-              envVars
+              envVars: JSON.stringify(currentEnvs)
             }
           }
         })
@@ -362,6 +392,34 @@ export class ShadowHandler {
     );
   }
 
+  public async updateSecureTunnelShadow(
+    secureTunnelShadowUpdate: SecureTunnelShadowUpdate
+  ) {
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'secure-tunnel', 'update'),
+      JSON.stringify(
+        buildUpdateShadowMessage({
+          reported: secureTunnelShadowUpdate,
+          clientId: this.clientId
+        })
+      )
+    );
+  }
+
+  public async updateDeviceAgentStatusShadow(
+    deviceAgentStatusShadowUpdate: DeviceAgentStatusShadowUpdate
+  ) {
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'device-agent-status', 'update'),
+      JSON.stringify(
+        buildUpdateShadowMessage({
+          reported: deviceAgentStatusShadowUpdate,
+          clientId: this.clientId
+        })
+      )
+    );
+  }
+
   public getProjectShadowUpdates() {
     this.publisher.publish(
       getShadowTopic(this.clientId, 'projects', 'get'),
@@ -373,7 +431,7 @@ export class ShadowHandler {
     this.publisher.publish(
       getShadowTopic(this.clientId, 'projects', 'update'),
       JSON.stringify(
-        buildUpdateProjectShadowMessage({
+        buildUpdateShadowMessage({
           clientId: this.clientId,
           reported: { [projectId]: null },
           desired: { [projectId]: null }

package/src/cloud-connection/shadow.ts

@@ -1,6 +1,7 @@
 import { logger } from '../util/logger';
 import { readAppCfgFile } from '../application-control';
 import { AppConfig } from '@alwaysai/app-configuration-schemas';
+import { stringifyError } from 'alwaysai/lib/util';
 
 export type AppConfigModels = {
   [modelId: string]: number;
@@ -43,7 +44,9 @@ export const getAppCfgModelsDiff = async ({
         newScripts[scriptName] = shadowScripts[scriptName];
       });
   } catch (e) {
-    logger.error(`Error parsing app config update: ${e.message}`);
+    logger.error(
+      `Error parsing app config update for ${projectId}!\n${stringifyError(e)}`
+    );
   }
 
   return { scripts: newScripts, updatedModels, untouchedModels };

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

@@ -7,6 +7,7 @@ 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()
@@ -14,8 +15,8 @@ const mockClient = {
 const clientId = 'test-client';
 
 const mockLiveUpdatesHandler = {
-  enableTransactionStatus: jest.fn(),
-  disableTransactionStatus: jest.fn()
+  enable: jest.fn(),
+  disable: jest.fn()
 } as any as LiveUpdatesHandler;
 
 describe('Test Transaction Manager', () => {
@@ -187,7 +188,7 @@ describe('Test Transaction Manager', () => {
       throw new Error('Expected start transaction to fail!');
     } catch (e) {
       console.log(e);
-      expect(e.code).toBe(txnMgr.Errors.PROJECT_ONGOING);
+      expect(e.code).toBe(txnMgr.Errors.PROJECT_TRANSACTION_ONGOING);
     }
     expect(txnMgr.getTransactionFromProject(projectId)).toEqual(txId);
     expect(txnMgr.getProjectFromTransaction(txId)).toEqual(projectId);
@@ -245,4 +246,127 @@ describe('Test Transaction Manager', () => {
 
     expect(txnMgr.getProjectFromTransaction(txId)).toBeUndefined();
   });
+
+  test('add an appCfgUpdate from txDetails', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    const cfgUpdate: ShadowUpdate = { txId: txId, projectId: projectId };
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
+
+    expect(txnMgr.getAppCfgUpdateFromTxID(txId)).toEqual(undefined);
+
+    txnMgr.setAppCfgUpdateToTx(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.getAppCfgUpdateFromTxID(txId)).toEqual(cfgUpdate);
+  });
+
+  test('update a tx while holding appCfgUpdate', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    const cfgUpdate: ShadowUpdate = { txId: txId, projectId: projectId };
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
+
+    txnMgr.setAppCfgUpdateToTx(txId, cfgUpdate);
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: false,
+      stepName: 'step2'
+    });
+
+    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.getAppCfgUpdateFromTxID(txId)).toEqual(cfgUpdate);
+  });
+
+  test('update an appCfgUpdate to the same tx ', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    const cfgUpdate: ShadowUpdate = { txId: txId, projectId: projectId };
+
+    const ogAppCfg1: AppConfigUpdate = {
+      newAppCfg: {
+        scripts: {
+          start: 'python app.py'
+        },
+        models: {}
+      }
+    };
+
+    const cfgUpdate2: ShadowUpdate = {
+      txId: txId,
+      projectId: projectId,
+      appCfgUpdate: ogAppCfg1
+    };
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
+
+    txnMgr.setAppCfgUpdateToTx(txId, cfgUpdate);
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: false,
+      stepName: 'step2'
+    });
+    txnMgr.setAppCfgUpdateToTx(txId, cfgUpdate2);
+
+    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.getAppCfgUpdateFromTxID(txId)).toEqual(cfgUpdate2);
+  });
+
+  test('remove a tx, ensure appCfgUpdate is removed', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    const cfgUpdate: ShadowUpdate = { txId: txId, projectId: projectId };
+
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
+
+    txnMgr.setAppCfgUpdateToTx(txId, cfgUpdate);
+    txnMgr.completeTransaction(txId);
+
+    expect(txnMgr.getProjectFromTransaction(txId)).toEqual(undefined);
+    expect(txnMgr.getTransactionFromProject(projectId)).toEqual(undefined);
+    expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
+    expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
+    expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
+    expect(txnMgr.getAppCfgUpdateFromTxID(txId)).toEqual(undefined);
+  });
 });

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

@@ -1,13 +1,14 @@
 import {
-  StatusResponsePayload,
+  ToClientStatusResponseMessage,
   buildToClientStatusResponseMessage,
   keyMirrors
 } from '@alwaysai/device-agent-schemas';
 import { LiveUpdatesHandler } from './live-updates-handler';
 import { Publisher } from './publisher';
 import { logger } from '../util/logger';
-import { keyMirror } from 'alwaysai/lib/util';
+import { keyMirror, stringifyError } from 'alwaysai/lib/util';
 import { CodedError } from '@carnesen/coded-error';
+import { ShadowUpdate } from './shadow-handler';
 
 interface TransactionDetails {
   txId: string;
@@ -16,6 +17,7 @@ interface TransactionDetails {
   start: string;
   update?: string;
   stop?: string;
+  appCfgUpdate?: ShadowUpdate;
 }
 
 export class TransactionManager {
@@ -24,11 +26,12 @@ export class TransactionManager {
   private liveUpdatesHandler: LiveUpdatesHandler;
   private publisher: Publisher;
 
-  private startTransaction(
+  private async startTransaction(
     txId: string,
     projectId: string,
+    liveUpdatesPublishFn?: () => Promise<void>,
     stepName?: string
-  ): void {
+  ): Promise<void> {
     // Check if the transaction already exists
     if (this.detailsByTx[txId]) {
       const txnDetails = this.detailsByTx[txId];
@@ -51,7 +54,7 @@ export class TransactionManager {
           null,
           2
         )}`,
-        this.Errors.PROJECT_ONGOING
+        this.Errors.PROJECT_TRANSACTION_ONGOING
       );
     }
 
@@ -65,10 +68,14 @@ export class TransactionManager {
     this.detailsByTx[txId] = txDetails;
     this.detailsByProject[projectId] = txDetails;
     logger.info(`Started transaction:\n${JSON.stringify(txDetails, null, 2)}`);
-    // send live updates
-    void this.liveUpdatesHandler.enableTransactionStatus({
-      txId
-    });
+
+    if (liveUpdatesPublishFn) {
+      await this.liveUpdatesHandler.enable(
+        keyMirrors.toClientMessageType.status_response,
+        liveUpdatesPublishFn,
+        txId
+      );
+    }
   }
 
   private updateTransaction(txId: string, stepName?: string): void {
@@ -93,7 +100,7 @@ export class TransactionManager {
 
   public Errors = keyMirror({
     TRANSACTION_ONGOING: null,
-    PROJECT_ONGOING: null,
+    PROJECT_TRANSACTION_ONGOING: null,
     TRANSACTION_NOT_ONGOING: null
   });
 
@@ -102,48 +109,49 @@ export class TransactionManager {
     projectId: string;
     txId: string;
     start: boolean;
+    liveUpdatesPublishFn?: () => Promise<void>;
     stepName?: string;
   }) {
-    const { func, projectId, txId, start, stepName } = props;
+    const { func, projectId, txId, start, liveUpdatesPublishFn, stepName } =
+      props;
     if (start) {
-      this.startTransaction(txId, projectId, stepName);
+      await this.startTransaction(
+        txId,
+        projectId,
+        liveUpdatesPublishFn,
+        stepName
+      );
     } else {
       this.updateTransaction(txId, stepName);
     }
     try {
       const completed = await func();
       if (completed) {
-        this.completeTransaction(txId);
-        const successStatusResponsePayload: StatusResponsePayload = {
-          status: keyMirrors.statusResponse.success
-        };
-        // Send final status message
-        const message = buildToClientStatusResponseMessage(
-          this.publisher.getClientId(),
-          successStatusResponsePayload,
-          txId
+        this.completeTransaction(
+          txId,
+          buildToClientStatusResponseMessage(
+            this.publisher.getClientId(),
+            { status: keyMirrors.statusResponse.success },
+            txId
+          )
         );
-        this.publisher.publishToClient(message);
       }
     } catch (e) {
-      const message: string = e.message;
       logger.error(
-        `Failed to execute cmd for ${projectId}:\n${message}\n${e.stack}`
+        `Failed to execute cmd for ${projectId}!\n${stringifyError(e)}`
       );
 
-      this.completeTransaction(txId);
-
-      const failureStatusResponsePayload: StatusResponsePayload = {
-        status: keyMirrors.statusResponse.failure,
-        message
-      };
-      // Send final status message
-      const failureStatusResponseMessage = buildToClientStatusResponseMessage(
-        this.publisher.getClientId(),
-        failureStatusResponsePayload,
-        txId
+      this.completeTransaction(
+        txId,
+        buildToClientStatusResponseMessage(
+          this.publisher.getClientId(),
+          {
+            status: keyMirrors.statusResponse.failure,
+            message: e.message
+          },
+          txId
+        )
       );
-      this.publisher.publishToClient(failureStatusResponseMessage);
     }
   }
 
@@ -169,7 +177,23 @@ export class TransactionManager {
     return txnDetails?.projectId;
   }
 
-  public completeTransaction(txId: string): void {
+  public getAppCfgUpdateFromTxID(txId: string): ShadowUpdate | undefined {
+    return this.detailsByTx[txId]?.appCfgUpdate;
+  }
+
+  public setAppCfgUpdateToTx(txId: string, appCfgUpdate: ShadowUpdate) {
+    if (this.isOngoingTransaction(txId)) {
+      this.detailsByTx[txId].appCfgUpdate = appCfgUpdate;
+    } else
+      throw new Error(
+        `Could not set appCfgUpdate, the transaction ${txId} does not exist.`
+      );
+  }
+
+  public completeTransaction(
+    txId: string,
+    messageToPublish?: ToClientStatusResponseMessage
+  ): void {
     const txDetails = this.detailsByTx[txId];
     if (txDetails === undefined) {
       throw new CodedError(
@@ -184,8 +208,13 @@ export class TransactionManager {
     delete this.detailsByTx[txId];
     delete this.detailsByProject[txDetails.projectId];
 
-    void this.liveUpdatesHandler.disableTransactionStatus({
+    this.liveUpdatesHandler.disable(
+      keyMirrors.toClientMessageType.status_response,
       txId
-    });
+    );
+
+    if (messageToPublish) {
+      this.publisher.publishToClient(messageToPublish);
+    }
   }
 }