@alwaysai/device-agent 1.5.0 → 2.0.1

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

@@ -68,13 +68,11 @@ describe('Test Shadow Handler', () => {
 
   test('handle project shadow empty delta', async () => {
     const payload = {
-      desired: {
-        [projectId1]: {}
-      }
+      [projectId1]: {}
     };
 
     const updates = await shadowHandler.handleProjectShadow({
-      topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+      topic: getShadowTopic(clientId, 'projects', 'update/delta'),
       payload,
       clientToken: ''
     });
@@ -104,16 +102,14 @@ describe('Test Shadow Handler', () => {
 
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              appConfig: JSON.stringify(appCfg1)
-            }
+          [projectId1]: {
+            appConfig: JSON.stringify(appCfg1)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -153,16 +149,14 @@ describe('Test Shadow Handler', () => {
 
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              appConfig: JSON.stringify(appCfg1)
-            }
+          [projectId1]: {
+            appConfig: JSON.stringify(appCfg1)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -220,19 +214,17 @@ describe('Test Shadow Handler', () => {
       };
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              appConfig: JSON.stringify(appCfg1)
-            },
-            [projectId2]: {
-              appConfig: JSON.stringify(appCfg2)
-            }
+          [projectId1]: {
+            appConfig: JSON.stringify(appCfg1)
+          },
+          [projectId2]: {
+            appConfig: JSON.stringify(appCfg2)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -281,16 +273,14 @@ describe('Test Shadow Handler', () => {
 
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              appConfig: JSON.stringify(appCfg1)
-            }
+          [projectId1]: {
+            appConfig: JSON.stringify(appCfg1)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -365,10 +355,8 @@ describe('Test Shadow Handler', () => {
 
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              appConfig: appCfg1 // This is missing JSON.stringify() making this an unparsable object.
-            }
+          [projectId1]: {
+            appConfig: appCfg1 // This is missing JSON.stringify() making this an unparsable object.
           }
         }
       };
@@ -378,7 +366,7 @@ describe('Test Shadow Handler', () => {
         .mockReturnValue({} as unknown as Logger);
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -393,13 +381,15 @@ describe('Test Shadow Handler', () => {
   describe('handle project shadow env vars', () => {
     test('handle a response from the getAccepted from the cloud', async () => {
       const envVars1 = {
-        VAR0: 'value0'
+        service: {
+          VAR0: 'value0'
+        }
       };
       const payload = {
         state: {
           reported: {
             [projectId1]: {
-              envVars: { service: envVars1 }
+              envVars: JSON.stringify(envVars1)
             }
           }
         }
@@ -414,22 +404,20 @@ describe('Test Shadow Handler', () => {
     });
     test('handle project shadow env vars update delta', async () => {
       const envVars1 = {
-        VAR1: 'value1'
+        service: {
+          VAR1: 'value1'
+        }
       };
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              envVars: {
-                service: envVars1
-              }
-            }
+          [projectId1]: {
+            envVars: JSON.stringify(envVars1)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -438,36 +426,36 @@ describe('Test Shadow Handler', () => {
         projectId: projectId1,
         txId: expect.any(String),
         envVarUpdate: {
-          envVars: { service: envVars1 }
+          envVars: envVars1
         }
       });
     });
 
     test('handle project shadow env vars update delta for two projects', async () => {
       const service1 = {
-        VAR1: 'value1'
+        service: {
+          VAR1: 'value1'
+        }
       };
       const service2 = {
-        VAR2: 'value2'
+        service: {
+          VAR2: 'value2'
+        }
       };
 
       const payload = {
         state: {
-          desired: {
-            [projectId1]: {
-              envVars: {
-                service: service1
-              }
-            },
-            [projectId2]: {
-              envVars: { service: service2 }
-            }
+          [projectId1]: {
+            envVars: JSON.stringify(service1)
+          },
+          [projectId2]: {
+            envVars: JSON.stringify(service2)
           }
         }
       };
 
       const updates = await shadowHandler.handleProjectShadow({
-        topic: getShadowTopic(clientId, 'projects', 'update/accepted'),
+        topic: getShadowTopic(clientId, 'projects', 'update/delta'),
         payload,
         clientToken: ''
       });
@@ -476,14 +464,14 @@ describe('Test Shadow Handler', () => {
         projectId: projectId1,
         txId: expect.any(String),
         envVarUpdate: {
-          envVars: { service: service1 }
+          envVars: service1
         }
       });
       expect(updates[1]).toEqual({
         projectId: projectId2,
         txId: expect.any(String),
         envVarUpdate: {
-          envVars: { service: service2 }
+          envVars: service2
         }
       });
     });

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

@@ -3,28 +3,39 @@ 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';
+import {
+  buildBaseShadowMessage,
+  DeviceAgentStatusShadowUpdate,
+  ShadowProjectsUpdateAll
+} from '@alwaysai/device-agent-schemas/lib/shadow-schema';
+import { stringifyError } from 'alwaysai/lib/util';
+import {
   getAllEnvs,
   readAppCfgFile,
-  setEnv
+  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 {
-  generateTxId,
-  validateProjectShadowUpdate,
-  buildBaseShadowMessage,
-  buildUpdateProjectShadowMessage,
-  buildUpdateSystemInfoShadowMessage,
-  getShadowTopic,
-  ShadowProjectsUpdateAll,
-  getDesiredFromMessage,
-  ProjectShadowUpdate,
-  buildUpdateSecureTunnelShadowMessage,
-  SecureTunnelShadowDescriptionReported
-} from '@alwaysai/device-agent-schemas';
+import { pruneModels } from '../application-control/models';
+import { MessageHandler } from './message-dispatcher';
+import { BaseHandler } from './base-message-handler';
 
 export type AppConfigUpdate = {
   newAppCfg: AppConfig;
@@ -112,9 +123,11 @@ export class ShadowHandler {
     // 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;
     }
@@ -145,6 +158,40 @@ export class ShadowHandler {
     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
   }: {
@@ -218,7 +265,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
@@ -242,24 +296,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) {
-          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) {
+      case this.shadowTopics.project.updateDelta: {
+        const delta = getUpdateDeltaStateFromMessage(payload);
+        if (!delta) {
           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;
       }
@@ -274,21 +318,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;
@@ -310,7 +351,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
+        })
       )
     );
   }
@@ -322,13 +366,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
         })
@@ -344,15 +388,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)
             }
           }
         })
@@ -361,15 +405,29 @@ export class ShadowHandler {
   }
 
   public async updateSecureTunnelShadow(
-    secureTunnelShadowUpdate: SecureTunnelShadowDescriptionReported
+    secureTunnelShadowUpdate: SecureTunnelShadowUpdate
   ) {
     this.publisher.publish(
       getShadowTopic(this.clientId, 'secure-tunnel', 'update'),
       JSON.stringify(
-        buildUpdateSecureTunnelShadowMessage(
-          secureTunnelShadowUpdate,
-          this.clientId
-        )
+        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
+        })
       )
     );
   }
@@ -385,7 +443,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 }
@@ -394,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/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 };