@alwaysai/device-agent 2.0.0 → 2.0.1

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/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;
@@ -440,3 +452,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,13 +1,17 @@
 import {
+  EnvVars,
   ToClientStatusResponseMessage,
   generateTxId,
-  keyMirrors
+  keyMirrors,
+  buildToClientStatusResponseMessage
 } 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';
+import { AppContent } from './device-agent-message-handler';
+import { AppConfig } from '@alwaysai/app-configuration-schemas';
 
 const mockClient = {
   publish: jest.fn()
@@ -58,11 +62,22 @@ describe('Test Transaction Manager', () => {
   test('Start a new transaction which completes in one step', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
+
+    const successFn = (txId: string) => {
+      const msg = buildToClientStatusResponseMessage(
+        clientId,
+        { status: keyMirrors.statusResponse.success },
+        txId
+      );
+      publisher.publishToClient(msg);
+    };
+
     await txnMgr.runTransactionStep({
       func: func_complete,
       projectId,
       txId,
-      start: true
+      start: true,
+      successFn
     });
     expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
     expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
@@ -198,6 +213,19 @@ describe('Test Transaction Manager', () => {
   test('Handle error in step function', async () => {
     const txId = generateTxId();
     const projectId = generateRandomProjectId();
+
+    const errorFn = (txId: string, errorMsg: string) => {
+      const msg = buildToClientStatusResponseMessage(
+        clientId,
+        {
+          status: keyMirrors.statusResponse.failure,
+          message: errorMsg
+        },
+        txId
+      );
+      publisher.publishToClient(msg);
+    };
+
     await txnMgr.runTransactionStep({
       func: jest.fn().mockImplementation(() => {
         throw new Error('Test error!');
@@ -205,7 +233,8 @@ describe('Test Transaction Manager', () => {
       projectId,
       txId,
       start: true,
-      stepName: 'step1'
+      stepName: 'step1',
+      errorFn
     });
     expect(txnMgr.isOngoingTransaction(txId)).toBe(false);
     expect(txnMgr.isOngoingTransactionForProjectID(projectId)).toBe(false);
@@ -369,4 +398,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

@@ -9,6 +9,7 @@ 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 +19,18 @@ 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 readonly liveUpdatesHandler: LiveUpdatesHandler;
   private publisher: Publisher;
-
   private async startTransaction(
     txId: string,
     projectId: string,
@@ -111,9 +116,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 +142,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 +185,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 +198,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 +230,5 @@ export class TransactionManager {
       keyMirrors.toClientMessageType.status_response,
       txId
     );
-
-    if (messageToPublish) {
-      this.publisher.publishToClient(messageToPublish);
-    }
   }
 }

package/src/jobs/job-handler.ts

@@ -0,0 +1,146 @@
+import { MessageHandler } from '../cloud-connection/message-dispatcher';
+import { logger } from '../util/logger';
+import {
+  BaseHandler,
+  HandlerContext
+} from '../cloud-connection/base-message-handler';
+import { DeviceAgentMessageHandler } from '../cloud-connection/device-agent-message-handler';
+
+class JobState {
+  private static instance: JobState;
+  public jobInProgress = false;
+
+  // Singleton pattern
+  private constructor() { } // eslint-disable-line
+
+  static getInstance(): JobState {
+    if (!JobState.instance) {
+      JobState.instance = new JobState();
+    }
+    return JobState.instance;
+  }
+}
+
+export class JobHandler extends BaseHandler implements MessageHandler {
+  private readonly msgHandler: DeviceAgentMessageHandler;
+
+  constructor(handlerContext: HandlerContext) {
+    super(handlerContext);
+    this.msgHandler = new DeviceAgentMessageHandler(
+      handlerContext,
+      this.handleJobError,
+      this.handleJobSuccess
+    );
+  }
+  private readonly state = JobState.getInstance();
+  private jobId = '';
+
+  public async handle(message: any, topic: string): Promise<void> {
+    const TOPICS = this.getJobTopic();
+
+    try {
+      switch (topic) {
+        case TOPICS.NOTIFY_NEXT:
+          await this.notifyNext(message);
+          break;
+        case TOPICS.START_NEXT_ACCEPTED:
+          this.startNextAccepted(message);
+          break;
+        default:
+          logger.info(`No handler for topic: ${topic}`);
+      }
+    } catch (error) {
+      logger.error(`Error handling job for topic: ${topic}`, error);
+    }
+  }
+
+  getJobTopic() {
+    const TOPICS = {
+      NOTIFY_NEXT: `$aws/things/${this.clientId}/jobs/notify-next`,
+      START_NEXT: `$aws/things/${this.clientId}/jobs/start-next`,
+      START_NEXT_ACCEPTED: `$aws/things/${this.clientId}/jobs/start-next/accepted`,
+      UPDATE: `$aws/things/${this.clientId}/jobs/${this.jobId}/update`
+    };
+    return TOPICS;
+  }
+
+  private async notifyNext(message: any): Promise<void> {
+    // Triggers when a job is completed and we need to notify the next job
+    logger.info('Invoking NOTIFY_NEXT handler');
+    const TOPICS = this.getJobTopic();
+    this.publisher.publish(TOPICS.START_NEXT, JSON.stringify({}));
+  }
+
+  private startNextAccepted(message: any): void {
+    if (this.state.jobInProgress) {
+      logger.info('Job already in progress');
+      return;
+    }
+
+    this.state.jobInProgress = true;
+
+    try {
+      if (message.execution) {
+        this.processJobExecution(message.execution);
+      } else {
+        this.state.jobInProgress = false;
+      }
+    } catch (error: unknown) {
+      this.handleError(error);
+    }
+    this.state.jobInProgress = false;
+  }
+
+  private processJobExecution(execution: any): void {
+    logger.info('Job in progress...');
+    const { jobDocument, jobId } = execution;
+    this.jobId = jobId;
+    this.msgHandler.handle(jobDocument);
+  }
+
+  private handleError(error: unknown): void {
+    if (error instanceof Error) {
+      logger.error(`Error starting next job`, error);
+      this.handleJobError('', error.message);
+    } else {
+      let errorMessage: string;
+      if (typeof error === 'string') {
+        errorMessage = error;
+      } else if (typeof error === 'object' && error !== null) {
+        try {
+          errorMessage = JSON.stringify(error);
+        } catch (jsonError) {
+          errorMessage = 'Unknown error (failed to stringify)';
+        }
+      } else {
+        errorMessage = 'Unknown error';
+      }
+      logger.error(`Unexpected error: ${errorMessage}`);
+      this.handleJobError('', errorMessage);
+    }
+  }
+
+  private updateJobStatus(status: string, message?: string): void {
+    const { UPDATE } = this.getJobTopic();
+    const payload = {
+      status,
+      statusDetails: {
+        message
+      }
+    };
+
+    this.publisher.publish(UPDATE, JSON.stringify(payload));
+    logger.info(`Marked job ${this.jobId} as ${status}`);
+    // Reset the job state
+    this.state.jobInProgress = false;
+  }
+
+  private readonly handleJobSuccess = (txId: string): void => {
+    this.updateJobStatus('SUCCEEDED');
+  };
+
+  private readonly handleJobError = (_txId: string, msg: string): void => {
+    logger.error(`Unexpected error type: ${msg}`);
+    this.updateJobStatus('FAILED', `Unexpected error: ${msg}`);
+  };
+}

package/src/secure-tunneling/secure-tunnel-message-handler.ts

@@ -0,0 +1,56 @@
+import {
+  getUpdateDeltaStateFromMessage,
+  validateSecureTunnelShadowUpdate
+} from '@alwaysai/device-agent-schemas';
+import { MessageHandler } from '../cloud-connection/message-dispatcher';
+import { logger } from '../util/logger';
+import { BaseHandler } from '../cloud-connection/base-message-handler';
+
+export class SecureTunnelMessageHandler
+  extends BaseHandler
+  implements MessageHandler<any>
+{
+  public getNotifyTopic(): string {
+    return `$aws/things/${this.clientId}/tunnels/notify`;
+  }
+
+  public async handle(message: any, topic: string): Promise<void> {
+    const secureTunnelNotifyTopic = this.getNotifyTopic();
+    if (topic === secureTunnelNotifyTopic) {
+      await this.secureTunnelHandler.secureTunnelNotifyHandler(message);
+    } else if (
+      topic === this.shadowHandler.shadowTopics.secureTunnel.updateDelta
+    ) {
+      logger.info(`Received secure tunnel update: ${JSON.stringify(message)}`);
+      await this.handleSecureTunnelMessage(message);
+    } else if (
+      topic === this.shadowHandler.shadowTopics.secureTunnel.deleteAccepted
+    ) {
+      logger.info(`Received secure tunnel deleteAccepted: ${message}`);
+      await this.secureTunnelHandler.destroy();
+    }
+  }
+
+  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);
+  }
+}

package/src/subcommands/app/version.ts

@@ -37,10 +37,26 @@ export const installAppCliLeaf = CliLeaf({
       description: 'Release Hash',
       required: false,
       hidden: true
+    }),
+    appCfg: CliStringInput({
+      description: 'Application Configuration',
+      required: false,
+      hidden: true
+    }),
+    envVars: CliStringInput({
+      description: 'Environment Variables configuration',
+      hidden: true,
+      required: false
     })
   },
   async action(_, opts) {
-    const { project, releaseHash, 'release-hash': releaseHashNew } = opts;
+    const {
+      project,
+      releaseHash,
+      'release-hash': releaseHashNew,
+      appCfg,
+      envVars
+    } = opts;
     if (releaseHash) {
       logger.warn(
         `--releaseHash is deprecated and will be removed in a future release. Please switch to --release-hash`
@@ -60,7 +76,9 @@ export const installAppCliLeaf = CliLeaf({
       payload: {
         baseCommand: keyMirrors.appVersionControl.install,
         projectId: project,
-        appReleaseHash: releaseHashResolved
+        appReleaseHash: releaseHashResolved,
+        appCfg,
+        envVars
       }
     };
     await deviceAgent.handleMessage(topic, message);