@alwaysai/device-agent 1.3.0 → 1.3.1

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

@@ -1,12 +1,36 @@
-import { generateTxId } from '@alwaysai/device-agent-schemas';
+import {
+  ToClientStatusResponseMessage,
+  generateTxId,
+  keyMirrors
+} 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';
+
+const mockClient = {
+  publish: jest.fn()
+};
+const clientId = 'test-client';
+
+const mockLiveUpdatesHandler = {
+  enableTransactionStatus: jest.fn(),
+  disableTransactionStatus: jest.fn()
+} as any as LiveUpdatesHandler;
 
 describe('Test Transaction Manager', () => {
   let txnMgr: TransactionManager;
+  let publisher: Publisher;
+  const func_complete: () => Promise<boolean> = jest
+    .fn()
+    .mockResolvedValue(true);
+  const func_incomplete: () => Promise<boolean> = jest
+    .fn()
+    .mockResolvedValue(false);
 
   beforeEach(() => {
-    txnMgr = new TransactionManager();
+    publisher = new Publisher(mockClient, clientId);
+    txnMgr = new TransactionManager(publisher, mockLiveUpdatesHandler);
     jest.clearAllMocks();
   });
 
@@ -14,60 +38,211 @@ describe('Test Transaction Manager', () => {
     return uuidv4();
   }
 
-  test('Add a new transaction', async () => {
+  test('Start a new transaction', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
-    txnMgr.addTransaction(txId, projectId);
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true
+    });
     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);
+  });
+
+  test('Start a new transaction which completes in one step', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    await txnMgr.runTransactionStep({
+      func: func_complete,
+      projectId,
+      txId,
+      start: true
+    });
+    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);
+  });
+
+  test('Start a new transaction and complete in second step', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true,
+      stepName: 'step1'
+    });
+    await txnMgr.runTransactionStep({
+      func: func_complete,
+      projectId,
+      txId,
+      start: false,
+      stepName: 'step2'
+    });
+    expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
+    expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
+    expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
   });
 
-  test('Add multiple transactions to different projects', async () => {
+  test('Start multiple transactions for different projects', async () => {
     const numTransactions = 3;
     const projectTxIdList: any = [];
     for (let i = 0; i < numTransactions; i++) {
       const txId = generateTxId();
       const projectId = generateRandomProjectId();
-      txnMgr.addTransaction(txId, projectId);
+      await txnMgr.runTransactionStep({
+        func: func_incomplete,
+        projectId,
+        txId,
+        start: true
+      });
       projectTxIdList.push({ projectId, txId });
     }
     projectTxIdList.forEach(({ txId, projectId }) => {
       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);
     });
   });
 
-  test('Attempt to add existing transaction to a project, results in failure', async () => {
+  test('Attempt to start an ongoing transaction, results in failure', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
-    txnMgr.addTransaction(txId, projectId);
-    expect(() => txnMgr.addTransaction(txId, projectId)).toThrow(
-      `Transaction ID ${txId} already exists for this projectId ${projectId}!`
-    );
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true,
+      stepName: 'step1'
+    });
+    try {
+      await txnMgr.runTransactionStep({
+        func: func_incomplete,
+        projectId,
+        txId,
+        start: true,
+        stepName: 'step2'
+      });
+      throw new Error('Expected starting transaction to fail!');
+    } catch (e) {
+      console.log(e);
+      expect(e.code).toBe(txnMgr.Errors.TRANSACTION_ONGOING);
+    }
     expect(txnMgr.getTransactionFromProject(projectId)).toEqual(txId);
     expect(txnMgr.getProjectFromTransaction(txId)).toEqual(projectId);
   });
 
-  test('Attempt to add a transaction to a project with an ongoing transaction, results in failure', async () => {
+  test('Attempt to continue a transaction that is not ongoing', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    try {
+      await txnMgr.runTransactionStep({
+        func: func_incomplete,
+        projectId,
+        txId,
+        start: false,
+        stepName: 'step1'
+      });
+      throw new Error('Expected continue transaction to fail!');
+    } catch (e) {
+      console.log(e);
+      expect(e.code).toBe(txnMgr.Errors.TRANSACTION_NOT_ONGOING);
+    }
+    expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
+    expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
+    expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
+  });
+
+  test('Attempt to start a transaction for a project with an ongoing transaction, results in failure', async () => {
     const txId = generateTxId();
     const txId2 = generateTxId();
     const projectId = generateRandomProjectId();
-    txnMgr.addTransaction(txId, projectId);
-    expect(() => txnMgr.addTransaction(txId2, projectId)).toThrow(
-      `This Project ID ${projectId} already has an ongoing transaction!`
-    );
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true,
+      stepName: 'step1'
+    });
+    try {
+      await txnMgr.runTransactionStep({
+        func: func_incomplete,
+        projectId,
+        txId: txId2,
+        start: true,
+        stepName: 'step2'
+      });
+      throw new Error('Expected start transaction to fail!');
+    } catch (e) {
+      console.log(e);
+      expect(e.code).toBe(txnMgr.Errors.PROJECT_ONGOING);
+    }
     expect(txnMgr.getTransactionFromProject(projectId)).toEqual(txId);
     expect(txnMgr.getProjectFromTransaction(txId)).toEqual(projectId);
     expect(txnMgr.getProjectFromTransaction(txId2)).toBeUndefined();
   });
 
-  test('Test remove transaction from queue', async () => {
+  test('Handle error in step function', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
-    txnMgr.addTransaction(txId, projectId);
+    await txnMgr.runTransactionStep({
+      func: jest.fn().mockImplementation(() => {
+        throw new Error('Test error!');
+      }),
+      projectId,
+      txId,
+      start: true,
+      stepName: 'step1'
+    });
+    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);
+  });
+
+  test('Complete ongoing transaction', async () => {
+    const txId = generateTxId();
+    const projectId = generateRandomProjectId();
+    await txnMgr.runTransactionStep({
+      func: func_incomplete,
+      projectId,
+      txId,
+      start: true,
+      stepName: 'step1'
+    });
     txnMgr.completeTransaction(txId);
 
     expect(txnMgr.getTransactionFromProject(projectId)).toBeUndefined();
     expect(txnMgr.getProjectFromTransaction(txId)).toBeUndefined();
+    expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
+    expect(txnMgr.isAnyOngoingTransaction()).toBe(false);
+  });
+
+  test('Test remove non-existing transaction from queue', async () => {
+    const txId = generateTxId();
+    try {
+      txnMgr.completeTransaction(txId);
+      throw new Error('Expected completTransaction to fail!');
+    } catch (e) {
+      console.log(e);
+      expect(e.code).toBe(txnMgr.Errors.TRANSACTION_NOT_ONGOING);
+    }
+
+    expect(txnMgr.getProjectFromTransaction(txId)).toBeUndefined();
   });
 });

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

@@ -1,43 +1,191 @@
+import {
+  StatusResponsePayload,
+  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 { CodedError } from '@carnesen/coded-error';
+
+interface TransactionDetails {
+  txId: string;
+  projectId: string;
+  stepName?: string;
+  start: string;
+  update?: string;
+  stop?: string;
+}
+
 export class TransactionManager {
-  private txToProject: Record<string, string> = {};
-  private projectToTx: Record<string, string> = {};
+  private detailsByTx: Record<string, TransactionDetails> = {};
+  private detailsByProject: Record<string, TransactionDetails> = {};
+  private liveUpdatesHandler: LiveUpdatesHandler;
+  private publisher: Publisher;
+
+  private startTransaction(
+    txId: string,
+    projectId: string,
+    stepName?: string
+  ): void {
+    // Check if the transaction already exists
+    if (this.detailsByTx[txId]) {
+      const txnDetails = this.detailsByTx[txId];
+      throw new CodedError(
+        `Transaction ${txId} already ongoing!\n${JSON.stringify(
+          txnDetails,
+          null,
+          2
+        )}`,
+        this.Errors.TRANSACTION_ONGOING
+      );
+    }
 
-  constructor() {
-    this.txToProject = {};
-    this.projectToTx = {};
+    // Check if there is any ongoing transactions for project
+    if (this.detailsByProject[projectId]) {
+      const txnDetails = this.detailsByProject[projectId];
+      throw new CodedError(
+        `Project ${projectId} already has an ongoing transaction!\n${JSON.stringify(
+          txnDetails,
+          null,
+          2
+        )}`,
+        this.Errors.PROJECT_ONGOING
+      );
+    }
+
+    // Map the Transaction ID with Project ID
+    const txDetails = {
+      txId,
+      projectId,
+      stepName,
+      start: new Date().toLocaleString()
+    };
+    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
+    });
   }
 
-  public addTransaction(txId: string, projectId: string) {
-    // Check if the Transaction already exists
-    if (this.txToProject[txId]) {
-      throw new Error(
-        `Transaction ID ${txId} already exists for this projectId ${projectId}!`
+  private updateTransaction(txId: string, stepName?: string): void {
+    const txDetails = this.detailsByTx[txId];
+    if (!txDetails) {
+      throw new CodedError(
+        `Cannot update transaction ${txId} since it doesn't exist!`,
+        this.Errors.TRANSACTION_NOT_ONGOING
       );
+    }
+    txDetails.stepName = stepName;
+    txDetails.update = new Date().toLocaleString();
+    logger.info(`Updated transaction:\n${JSON.stringify(txDetails, null, 2)}`);
+  }
+
+  constructor(publisher: Publisher, liveUpdatesHandler: LiveUpdatesHandler) {
+    this.detailsByTx = {};
+    this.detailsByProject = {};
+    this.publisher = publisher;
+    this.liveUpdatesHandler = liveUpdatesHandler;
+  }
+
+  public Errors = keyMirror({
+    TRANSACTION_ONGOING: null,
+    PROJECT_ONGOING: null,
+    TRANSACTION_NOT_ONGOING: null
+  });
+
+  public async runTransactionStep(props: {
+    func: () => Promise<boolean>;
+    projectId: string;
+    txId: string;
+    start: boolean;
+    stepName?: string;
+  }) {
+    const { func, projectId, txId, start, stepName } = props;
+    if (start) {
+      this.startTransaction(txId, projectId, stepName);
     } else {
-      // Check if there is any ongoing Transactions for Project ID
-      if (this.projectToTx[projectId]) {
-        throw new Error(
-          `This Project ID ${projectId} already has an ongoing transaction!`
+      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
         );
-      } else {
-        // Map the Transaction ID with Project ID
-        this.txToProject[txId] = projectId;
-        this.projectToTx[projectId] = 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}`
+      );
+
+      this.completeTransaction(txId);
+
+      const failureStatusResponsePayload: StatusResponsePayload = {
+        status: keyMirrors.statusResponse.failure,
+        message
+      };
+      // Send final status message
+      const failureStatusResponseMessage = buildToClientStatusResponseMessage(
+        this.publisher.getClientId(),
+        failureStatusResponsePayload,
+        txId
+      );
+      this.publisher.publishToClient(failureStatusResponseMessage);
     }
   }
 
-  public getTransactionFromProject(projectId: string) {
-    return this.projectToTx[projectId];
+  public getTransactionFromProject(projectId: string): string | undefined {
+    const txnDetails = this.detailsByProject[projectId];
+    return txnDetails?.txId;
+  }
+
+  public isOngoingTransactionForProjectID(projectId: string): boolean {
+    return projectId in this.detailsByProject;
   }
 
-  public getProjectFromTransaction(txId: string) {
-    return this.txToProject[txId];
+  public isOngoingTransaction(txId: string): boolean {
+    return txId in this.detailsByTx;
   }
 
-  public completeTransaction(txId: string) {
-    const projectId = this.txToProject[txId];
-    delete this.txToProject[txId];
-    delete this.projectToTx[projectId];
+  public isAnyOngoingTransaction(): boolean {
+    return Object.keys(this.detailsByTx).length > 0;
+  }
+
+  public getProjectFromTransaction(txId: string): string | undefined {
+    const txnDetails = this.detailsByTx[txId];
+    return txnDetails?.projectId;
+  }
+
+  public completeTransaction(txId: string): void {
+    const txDetails = this.detailsByTx[txId];
+    if (txDetails === undefined) {
+      throw new CodedError(
+        `Cannot complete transaction ${txId} since it doesn't exist!`,
+        this.Errors.TRANSACTION_NOT_ONGOING
+      );
+    }
+    txDetails.stop = new Date().toLocaleString();
+    logger.info(
+      `Completed transaction:\n${JSON.stringify(txDetails, null, 2)}`
+    );
+    delete this.detailsByTx[txId];
+    delete this.detailsByProject[txDetails.projectId];
+
+    void this.liveUpdatesHandler.disableTransactionStatus({
+      txId
+    });
   }
 }

package/src/subcommands/app/analytics.ts

@@ -0,0 +1,99 @@
+import {
+  CliFlagInput,
+  CliLeaf,
+  CliNumberInput,
+  CliStringInput
+} from '@alwaysai/alwayscli';
+import { readAppCfgFile } from '../../application-control';
+import { DeviceAgentCloudConnection } from '../../cloud-connection/device-agent-cloud-connection';
+import sleep from '../../util/sleep';
+import { logger } from '../../util/logger';
+import { assign, merge } from 'lodash';
+
+export const getAnalyticsCfgCliLeaf = CliLeaf({
+  name: 'get-analytics-cfg',
+  description: 'Get analytics configuration for an application',
+  namedInputs: {
+    project: CliStringInput({
+      description: 'Project Id',
+      required: true
+    })
+  },
+  async action(_, opts) {
+    const { project } = opts;
+    const appCfg = await readAppCfgFile({ projectId: project });
+    if (appCfg.analytics !== undefined) {
+      console.log(JSON.stringify(appCfg.analytics, null, 2));
+    } else {
+      console.log('No analytics configuration for app!');
+    }
+  }
+});
+
+export const setAnalyticsCfgCliLeaf = CliLeaf({
+  name: 'set-analytics-cfg',
+  description:
+    'Set analytics configuration for an application. Note that this resets the config so all desired options must be set',
+  namedInputs: {
+    project: CliStringInput({
+      description: 'Project Id',
+      required: true
+    }),
+    'enable-cloud-publish': CliFlagInput({
+      description: 'Enable publishing analytics to cloud'
+    }),
+    'enable-file-publish': CliFlagInput({
+      description: 'Enable publishing analytics to file'
+    }),
+    'file-size-bytes': CliNumberInput({
+      description: 'Set the max file size in bytes for analytics file writing',
+      required: false
+    })
+  },
+  async action(
+    _,
+    {
+      project,
+      'enable-cloud-publish': enableCLoudPublish,
+      'enable-file-publish': enableFilePublish,
+      'file-size-bytes': fileSizeBytes
+    }
+  ) {
+    const deviceAgent = new DeviceAgentCloudConnection();
+    await deviceAgent.setupHandlers();
+
+    const newAppCfg = {
+      analytics: {
+        enable_cloud_publish: enableCLoudPublish,
+        enable_file_publish: enableFilePublish,
+        file_size_bytes: fileSizeBytes
+      }
+    };
+    const existingAppCfg = await readAppCfgFile({ projectId: project });
+    const appCfg = assign(existingAppCfg, merge(existingAppCfg, newAppCfg));
+
+    // Update the shadow as a client
+    const topic = deviceAgent.getShadowTopics().projects.update;
+    const packet = {
+      state: {
+        desired: {
+          [project]: {
+            appConfig: JSON.stringify(appCfg) // Pack app config as string as dictated by schema
+          }
+        }
+      },
+      clientToken: 'client'
+    };
+    logger.debug(
+      `Publishing message:\n${JSON.stringify({ topic, packet }, null, 2)}`
+    );
+    deviceAgent.publisher.publish(topic, JSON.stringify(packet));
+    // Sleep for extra time to ensure time for shadow response
+    await sleep(10000);
+
+    while (deviceAgent.isCmdInProgress(project)) {
+      await sleep(1000);
+    }
+    await deviceAgent.stop();
+  }
+});

package/src/subcommands/app/index.ts

@@ -5,8 +5,7 @@ import {
   addModelCliLeaf,
   removeModelCliLeaf,
   replaceModelsCliLeaf,
-  updateModelsCliLeaf,
-  installModelCliLeaf
+  updateModelsCliLeaf
 } from './models';
 import {
   getAppStatusCliLeaf,
@@ -22,6 +21,7 @@ import {
   rollbackAppCliLeaf
 } from './version';
 import { getShadowCliLeaf, updateShadowCliLeaf } from './shadow';
+import { getAnalyticsCfgCliLeaf, setAnalyticsCfgCliLeaf } from './analytics';
 
 export const appCliBranch = CliBranch({
   name: 'app',
@@ -40,10 +40,11 @@ export const appCliBranch = CliBranch({
     addModelCliLeaf,
     removeModelCliLeaf,
     replaceModelsCliLeaf,
-    installModelCliLeaf,
     updateModelsCliLeaf,
     getAllEnvsCliLeaf,
     setEnvCliLeaf,
+    getAnalyticsCfgCliLeaf,
+    setAnalyticsCfgCliLeaf,
     getShadowCliLeaf,
     updateShadowCliLeaf
   ]

package/src/subcommands/app/models.ts

@@ -13,6 +13,7 @@ import {
 } from '../../application-control';
 import { DeviceAgentCloudConnection } from '../../cloud-connection/device-agent-cloud-connection';
 import sleep from '../../util/sleep';
+import { logger } from '../../util/logger';
 
 export const showAppModelsCliLeaf = CliLeaf({
   name: 'show-models',
@@ -52,66 +53,29 @@ export const addModelCliLeaf = CliLeaf({
     const deviceAgent = new DeviceAgentCloudConnection();
     await deviceAgent.setupHandlers();
 
-    const topic = deviceAgent.getShadowTopics().projects.updateDelta;
+    // Update the shadow as a client
+    const topic = deviceAgent.getShadowTopics().projects.update;
 
     const newAppCfg = await readAppCfgFile({ projectId: project });
     newAppCfg.models[model] = Number(version);
 
-    const message = {
-      version: 3,
-      timestamp: 0,
+    const packet = {
       state: {
-        [project]: {
-          appConfig: JSON.stringify(newAppCfg)
-        }
-      },
-      clientToken: 'not-self'
-    };
-
-    await deviceAgent.handleMessage(topic, message);
-    while (deviceAgent.isCmdInProgress(project)) {
-      await sleep(1000);
-    }
-    await deviceAgent.stop();
-  }
-});
-
-export const installModelCliLeaf = CliLeaf({
-  name: 'install-model',
-  description: 'Install an alwaysAI model to a project',
-  namedInputs: {
-    project: CliStringInput({
-      description: 'Project ID',
-      required: true
-    }),
-    modelName: CliStringInput({
-      description: 'Model Name',
-      required: true
-    }),
-    modelVersion: CliNumberInput({
-      description: 'Model Version',
-      required: true
-    })
-  },
-  async action(_, opts) {
-    const { project, modelName, modelVersion } = opts;
-    const deviceAgent = new DeviceAgentCloudConnection();
-    await deviceAgent.setupHandlers();
-    const topic = deviceAgent.getShadowTopics().projects.getAccepted;
-    const newAppCfg = await readAppCfgFile({ projectId: project });
-    newAppCfg['models'][modelName] = modelVersion;
-
-    const message = {
-      state: {
-        delta: {
+        desired: {
           [project]: {
             appConfig: JSON.stringify(newAppCfg)
           }
         }
       },
-      clientToken: deviceAgent.getClientId()
+      clientToken: 'client'
     };
-    await deviceAgent.handleMessage(topic, message);
+    logger.debug(
+      `Publishing message:\n${JSON.stringify({ topic, packet }, null, 2)}`
+    );
+    deviceAgent.publisher.publish(topic, JSON.stringify(packet));
+    // Sleep for extra time to ensure time for shadow response
+    await sleep(10000);
+
     while (deviceAgent.isCmdInProgress(project)) {
       await sleep(1000);
     }