@alwaysai/device-agent 1.3.1 → 1.5.0

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

@@ -2,28 +2,29 @@ import {
   AppConfig,
   validateAppConfig
 } from '@alwaysai/app-configuration-schemas';
-import { EnvVars, getAllEnvs, readAppCfgFile } from '../application-control';
+import {
+  EnvVars,
+  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 { getSystemInformation } from '../device-control/device-control';
-import { generateTxId } from '@alwaysai/device-agent-schemas';
-
-export interface ShadowTopics {
-  projects: {
-    get: string;
-    getAccepted: string;
-    getRejected: string;
-    update: string;
-    updateDelta: string;
-    updateAccepted: string;
-    updateRejected: string;
-    delete: string;
-  };
-  systemInfo: {
-    update: string;
-  };
-}
+import {
+  generateTxId,
+  validateProjectShadowUpdate,
+  buildBaseShadowMessage,
+  buildUpdateProjectShadowMessage,
+  buildUpdateSystemInfoShadowMessage,
+  getShadowTopic,
+  ShadowProjectsUpdateAll,
+  getDesiredFromMessage,
+  ProjectShadowUpdate,
+  buildUpdateSecureTunnelShadowMessage,
+  SecureTunnelShadowDescriptionReported
+} from '@alwaysai/device-agent-schemas';
 
 export type AppConfigUpdate = {
   newAppCfg: AppConfig;
@@ -44,101 +45,190 @@ export type ShadowUpdate = {
 export class ShadowHandler {
   private clientId: string;
   private publisher: Publisher;
-  public readonly shadowPrefix: string;
-  public readonly shadowTopics: ShadowTopics;
-
+  public projectShadowTopics: string[] = [];
+  public readonly shadowTopics: { [key: string]: any };
   constructor(clientId: string, publisher: Publisher) {
     this.clientId = clientId;
     this.publisher = publisher;
-    this.shadowPrefix = `$aws/things/${this.clientId}/shadow/name/`;
     this.shadowTopics = {
-      projects: {
-        get: `${this.shadowPrefix}projects/get`,
-        getAccepted: `${this.shadowPrefix}projects/get/accepted`,
-        getRejected: `${this.shadowPrefix}projects/get/rejected`,
-        update: `${this.shadowPrefix}projects/update`,
-        updateDelta: `${this.shadowPrefix}projects/update/delta`,
-        updateAccepted: `${this.shadowPrefix}projects/update/accepted`,
-        updateRejected: `${this.shadowPrefix}projects/update/rejected`,
-        delete: `${this.shadowPrefix}projects/delete`
+      project: {
+        get: getShadowTopic(clientId, 'projects', 'get'),
+        getAccepted: getShadowTopic(clientId, 'projects', 'get/accepted'),
+        getRejected: getShadowTopic(clientId, 'projects', 'get/rejected'),
+        update: getShadowTopic(clientId, 'projects', 'update'),
+        updateDelta: getShadowTopic(clientId, 'projects', 'update/delta'),
+        updateAccepted: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        updateRejected: getShadowTopic(clientId, 'projects', 'update/rejected')
       },
       systemInfo: {
-        update: `${this.shadowPrefix}system-info/update`
+        update: getShadowTopic(clientId, 'system-info', 'update')
+      },
+      secureTunnel: {
+        get: getShadowTopic(clientId, 'secure-tunnel', 'get'),
+        getAccepted: getShadowTopic(clientId, 'secure-tunnel', 'get/accepted'),
+        getRejected: getShadowTopic(clientId, 'secure-tunnel', 'get/rejected'),
+        update: getShadowTopic(clientId, 'secure-tunnel', 'update'),
+        updateDelta: getShadowTopic(clientId, 'secure-tunnel', 'update/delta'),
+        updateAccepted: getShadowTopic(
+          clientId,
+          'secure-tunnel',
+          'update/accepted'
+        ),
+        updateRejected: getShadowTopic(
+          clientId,
+          'secure-tunnel',
+          'update/rejected'
+        ),
+        delete: getShadowTopic(clientId, 'secure-tunnel', 'delete'),
+        deleteAccepted: getShadowTopic(
+          clientId,
+          'secure-tunnel',
+          'delete/accepted'
+        ),
+        deleteRejected: getShadowTopic(
+          clientId,
+          'secure-tunnel',
+          'delete/rejected'
+        )
       }
     };
+    this.projectShadowTopics.push(this.shadowTopics.project.getAccepted);
+    this.projectShadowTopics.push(this.shadowTopics.project.getRejected);
+    this.projectShadowTopics.push(this.shadowTopics.project.updateDelta);
+    this.projectShadowTopics.push(this.shadowTopics.project.updateAccepted);
+    this.projectShadowTopics.push(this.shadowTopics.project.updateRejected);
+  }
+
+  private async generateAppConfigUpdate({
+    appConfig,
+    txId,
+    projectId
+  }: {
+    appConfig: string;
+    txId: string;
+    projectId: string;
+  }): Promise<AppConfigUpdate | null> {
+    let newAppCfg: any;
+    // Handle errors and validation
+    try {
+      newAppCfg = JSON.parse(appConfig);
+    } catch (error) {
+      logger.error(
+        `Could not parse the appConfig for transaction ${txId}!\n${error}`
+      );
+      return null;
+    }
+
+    if (!validateAppConfig(newAppCfg)) {
+      // FIXME: Raise an exception to be handled at higher layer
+      logger.error(
+        `Received invalid app config for ${projectId}!\n${JSON.stringify(
+          validateAppConfig.errors,
+          null,
+          2
+        )}`
+      );
+      return null;
+    }
+
+    // If all ok, return the AppConfigUpdate
+    logger.info(`Found a delta for app config shadow. Updating ${projectId}`);
+    const { updatedModels } = await getAppCfgModelsDiff({
+      newAppCfg,
+      projectId
+    });
+
+    const appCfgUpdate: AppConfigUpdate = { newAppCfg };
+    if (updatedModels && Object.keys(updatedModels).length) {
+      appCfgUpdate.updatedModels = updatedModels;
+    }
+    return appCfgUpdate;
   }
 
-  private async handleNamedShadowUpdate({
+  private async processProjectShadowUpdates({
     delta
   }: {
     delta: any;
   }): Promise<ShadowUpdate[]> {
-    const updates: ShadowUpdate[] = [];
-
-    const deltaKeys = Object.keys(delta);
-
-    for (const projectId of deltaKeys) {
-      const projectShadow = delta[projectId];
-      // For incoming shadow updates, there will be no TxID, so it needs to be generated here.
-      const txId = generateTxId();
-      const shadowUpdate: ShadowUpdate = { projectId, txId };
+    const shadowUpdatesPromises: Promise<ShadowUpdate | null>[] = [];
 
-      if (projectShadow.appConfig) {
-        let newAppCfg: any;
-        try {
-          newAppCfg = JSON.parse(projectShadow.appConfig);
-        } catch (error) {
+    for (const [projectId, projectDelta] of Object.entries(delta)) {
+      if (projectDelta) {
+        const valid = validateProjectShadowUpdate(projectDelta);
+        if (!valid) {
           logger.error(
-            `Could not parse the appConfig for transaction ${txId}!\n${error}`
-          );
-          continue;
-        }
-        if (!validateAppConfig(newAppCfg)) {
-          // FIXME: Raise an exception to be handled at higher layer
-          logger.error(
-            `Received invalid app config for ${projectId}!\n${JSON.stringify(
-              validateAppConfig.errors,
+            `Error validating shadow update: ${JSON.stringify(
+              { projectDelta, errors: validateProjectShadowUpdate.errors },
               null,
               2
             )}`
           );
           continue;
         }
-        const { updatedModels } = await getAppCfgModelsDiff({
-          newAppCfg,
-          projectId
-        });
 
-        if (updatedModels && Object.keys(updatedModels).length) {
-          shadowUpdate.appCfgUpdate = { newAppCfg, updatedModels };
-        } else {
-          shadowUpdate.appCfgUpdate = { newAppCfg };
-        }
-      } else {
-        logger.info(
-          `Ignoring app config shadow update for ${projectId} due to no delta`
+        shadowUpdatesPromises.push(
+          this.processProjectShadowUpdate({
+            projectId,
+            projectDelta
+          })
         );
       }
+    }
+    return (await Promise.all(shadowUpdatesPromises)).filter(
+      (promise) => promise !== null
+    ) as ShadowUpdate[];
+  }
 
-      if (projectShadow.envVars) {
-        const envVars = projectShadow.envVars;
-        shadowUpdate.envVarUpdate = { envVars };
-      } else {
-        logger.info(
-          `Ignoring app environment variable shadow update for ${projectId} due to no delta`
-        );
-      }
-      if (shadowUpdate.appCfgUpdate || shadowUpdate.envVarUpdate) {
-        updates.push(shadowUpdate);
+  private async processProjectShadowUpdate({
+    projectId,
+    projectDelta
+  }: {
+    projectId: string;
+    projectDelta: ProjectShadowUpdate;
+  }): Promise<ShadowUpdate | null> {
+    const txId = generateTxId();
+    const shadowUpdate: ShadowUpdate = { projectId, txId };
+
+    // Handle appConfig Updates
+    if (!projectDelta.appConfig) {
+      logger.info(
+        `Ignoring app config shadow update for ${projectId} due to no delta`
+      );
+    } else {
+      logger.info(`Found a delta for app config shadow. Updating ${projectId}`);
+      const appConfig = projectDelta.appConfig;
+      const appCfgUpdate = await this.generateAppConfigUpdate({
+        appConfig,
+        txId,
+        projectId
+      });
+
+      if (appCfgUpdate) {
+        shadowUpdate.appCfgUpdate = appCfgUpdate;
       }
     }
 
-    return updates;
+    // Handle envVars Updates
+    if (!projectDelta.envVars) {
+      logger.info(
+        `Ignoring app environment variable shadow update for ${projectId} due to no delta`
+      );
+    } else {
+      logger.info(
+        `Found a delta for app environment variable shadow. Updating ${projectId}`
+      );
+      const envVars = projectDelta.envVars;
+      shadowUpdate.envVarUpdate = { envVars };
+    }
+
+    return shadowUpdate.appCfgUpdate || shadowUpdate.envVarUpdate
+      ? shadowUpdate
+      : null;
   }
 
   // Public interface
 
-  public async handleShadowTopic({
+  public async handleProjectShadow({
     topic,
     payload,
     clientToken
@@ -148,12 +238,14 @@ export class ShadowHandler {
     clientToken: string;
   }): Promise<ShadowUpdate[]> {
     // TODO: make use a function like the other topic getters
-    const shadowName = topic.split('/')[5];
+    if (!this.projectShadowTopics.includes(topic)) {
+      throw Error(`Topic ${topic} is not in the ${this.projectShadowTopics}`);
+    }
     switch (topic) {
-      case this.shadowTopics.projects.updateAccepted:
+      case this.shadowTopics.project.updateAccepted: {
         if (clientToken === this.clientId) {
           logger.debug(
-            `Ignoring delta caused by Device Agent: ${JSON.stringify(
+            `Ignoring message as it was caused by Device Agent itself: ${JSON.stringify(
               { topic, payload },
               null,
               2
@@ -161,19 +253,49 @@ export class ShadowHandler {
           );
           break;
         }
-        return await this.handleNamedShadowUpdate({ delta: payload.desired });
-      case this.shadowTopics.projects.getAccepted:
-        if (payload['delta']) {
-          return await this.handleNamedShadowUpdate({
-            delta: payload['delta']
+        const desired = getDesiredFromMessage(payload);
+        if (!desired) {
+          logger.debug(
+            `No desired state found in message: ${JSON.stringify(payload)}`
+          );
+        } else {
+          return await this.processProjectShadowUpdates({ delta: desired });
+        }
+        break;
+      }
+      case this.shadowTopics.project.getAccepted: {
+        if (clientToken !== this.clientId) {
+          logger.debug(
+            `Ignoring message as it was caused by Device Agent itself: ${JSON.stringify(
+              { topic, payload },
+              null,
+              2
+            )}`
+          );
+          break;
+        }
+        if (payload.state.delta) {
+          return await this.processProjectShadowUpdates({
+            delta: payload.state.delta
           });
         } else {
-          logger.info(`No delta updates in named shadow '${shadowName}'`);
+          logger.info(
+            `No delta in projects.getAccepted in named shadow '${
+              topic.split('/')[5]
+            }'`
+          );
         }
         break;
+      }
+      case this.shadowTopics.project.getRejected:
+      case this.shadowTopics.project.updateDelta:
+      case this.shadowTopics.project.updateRejected: {
+        // Not handling these for now
+        break;
+      }
       default:
         logger.info(
-          `Ignoring shadow message: ${JSON.stringify(
+          `Did not match a correct topic. Ignoring shadow message: ${JSON.stringify(
             { topic, payload },
             null,
             2
@@ -185,58 +307,90 @@ export class ShadowHandler {
 
   public async updateSystemInfoShadow() {
     const systemInfo = await getSystemInformation();
-    const packet = {
-      state: {
-        reported: systemInfo
-      },
-      clientToken: this.clientId
-    };
-    const topic = this.shadowTopics.systemInfo.update;
-    this.publisher.publish(topic, JSON.stringify(packet));
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'system-info', 'update'),
+      JSON.stringify(
+        buildUpdateSystemInfoShadowMessage(this.clientId, systemInfo)
+      )
+    );
   }
 
   public async updateProjectShadow(projectId: string) {
     const appCfg = await readAppCfgFile({ projectId });
     const envVars = await getAllEnvs({ projectId });
 
-    const packet = {
-      state: {
-        reported: {
-          [projectId]: {
-            appConfig: JSON.stringify(appCfg),
-            envVars
-          }
-        }
-      },
-      clientToken: this.clientId
+    const toReport: ShadowProjectsUpdateAll = {
+      [projectId]: {
+        appConfig: JSON.stringify(appCfg),
+        envVars
+      }
     };
-    const topic = this.shadowTopics.projects.update;
-    this.publisher.publish(topic, JSON.stringify(packet));
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'projects', 'update'),
+      JSON.stringify(
+        buildUpdateProjectShadowMessage({
+          clientId: this.clientId,
+          reported: toReport
+        })
+      )
+    );
   }
 
-  public getShadowUpdates() {
-    const topic = this.shadowTopics.projects.get;
-    const packet = {
-      clientToken: this.clientId
-    };
-    this.publisher.publish(topic, JSON.stringify(packet));
+  public async updateProjectEnvVars({
+    projectId,
+    envVars
+  }: {
+    projectId: string;
+    envVars: EnvVars;
+  }) {
+    await setEnv({ projectId, envVars });
+
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'projects', 'update'),
+      JSON.stringify(
+        buildUpdateProjectShadowMessage({
+          clientId: this.clientId,
+          reported: {
+            [projectId]: {
+              envVars
+            }
+          }
+        })
+      )
+    );
   }
 
-  public clearAppConfig(projectId: string) {
-    const topic = this.shadowTopics.projects.update;
-    // TODO: We should actually send only desired and handle the delta
-    // to update reported
-    const packet = {
-      state: {
-        desired: {
-          [projectId]: null
-        },
-        reported: {
-          [projectId]: null
-        }
-      },
-      clientToken: this.clientId
-    };
-    this.publisher.publish(topic, JSON.stringify(packet));
+  public async updateSecureTunnelShadow(
+    secureTunnelShadowUpdate: SecureTunnelShadowDescriptionReported
+  ) {
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'secure-tunnel', 'update'),
+      JSON.stringify(
+        buildUpdateSecureTunnelShadowMessage(
+          secureTunnelShadowUpdate,
+          this.clientId
+        )
+      )
+    );
+  }
+
+  public getProjectShadowUpdates() {
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'projects', 'get'),
+      JSON.stringify(buildBaseShadowMessage(this.clientId))
+    );
+  }
+
+  public clearProjectShadow(projectId: string) {
+    this.publisher.publish(
+      getShadowTopic(this.clientId, 'projects', 'update'),
+      JSON.stringify(
+        buildUpdateProjectShadowMessage({
+          clientId: this.clientId,
+          reported: { [projectId]: null },
+          desired: { [projectId]: null }
+        })
+      )
+    );
   }
 }

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()
@@ -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

@@ -8,6 +8,7 @@ import { Publisher } from './publisher';
 import { logger } from '../util/logger';
 import { keyMirror } 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 {
@@ -169,6 +171,19 @@ export class TransactionManager {
     return txnDetails?.projectId;
   }
 
+  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): void {
     const txDetails = this.detailsByTx[txId];
     if (txDetails === undefined) {