feat: error handling while in automation execution

Author: FedericoAmura

packages/automation/src/lib/listeners/listener.ts

@@ -1,8 +1,11 @@
 import { EventEmitter } from 'events';
 
+import { onError } from '../types';
+
 export interface ListenerParams {
   start?: () => Promise<void>;
   stop?: () => Promise<void>;
+  onError?: onError;
 }
 
 /**
@@ -23,16 +26,23 @@ export class Listener<T = unknown> {
    */
   public stop: () => Promise<void>;
 
+  /**
+   * The error handling function to call when an error occurs.
+   */
+  public onError?: onError;
+
   /**
    * Constructor for the Listener class.
    * @param params The parameters object containing start and stop functions.
    */
   constructor({
     start = async () => {},
     stop = async () => {},
+    onError,
   }: ListenerParams = {}) {
     this.start = start;
     this.stop = stop;
+    this.onError = onError;
   }
 
   /**

packages/automation/src/lib/state-machine.ts

@@ -1,6 +1,10 @@
 import { ethers } from 'ethers';
 
-import { LIT_RPC } from '@lit-protocol/constants';
+import {
+  AutomationError,
+  UnknownError,
+  LIT_RPC,
+} from '@lit-protocol/constants';
 import { LitContracts } from '@lit-protocol/contracts-sdk';
 import { LitNodeClient } from '@lit-protocol/lit-node-client';
 
@@ -34,6 +38,7 @@ export type MachineStatus = 'running' | 'stopped';
  */
 export class StateMachine {
   private readonly debug;
+  private readonly onError?: (error: unknown, context?: string) => void;
   private context: MachineContext;
 
   private litNodeClient: LitNodeClient;
@@ -51,6 +56,7 @@ export class StateMachine {
   constructor(params: BaseStateMachineParams) {
     this.id = this.generateId();
     this.debug = params.debug ?? false;
+    this.onError = params.onError;
     this.context = new MachineContext(params.context);
 
     this.litNodeClient = params.litNodeClient;
@@ -62,6 +68,7 @@ export class StateMachine {
   static fromDefinition(machineConfig: StateMachineDefinition): StateMachine {
     const {
       debug = false,
+      onError,
       litNodeClient,
       litContracts = {},
       privateKey,
@@ -94,6 +101,7 @@ export class StateMachine {
 
     const stateMachine = new StateMachine({
       debug,
+      onError,
       litNodeClient: litNodeClientInstance,
       litContracts: litContractsInstance,
       privateKey,
@@ -360,16 +368,18 @@ export class StateMachine {
       transitionConfig.listeners = listeners;
       // Aggregate (AND) all listener checks to a single function result
       transitionConfig.check = async (values) => {
-        console.log(
-          `${transition.fromState} -> ${transition.toState} values`,
-          values
-        );
+        debug &&
+          console.log(
+            `${transition.fromState} -> ${transition.toState} values`,
+            values
+          );
         return Promise.all(checks.map((check) => check(values))).then(
           (results) => {
-            console.log(
-              `${transition.fromState} -> ${transition.toState} results`,
-              results
-            );
+            debug &&
+              console.log(
+                `${transition.fromState} -> ${transition.toState} results`,
+                results
+              );
             return results.every((result) => result);
           }
         );
@@ -421,10 +431,15 @@ export class StateMachine {
       await this.transitionTo(toState);
     };
 
+    const onTransitionError = async (error: unknown) => {
+      this.handleError(error, `Error at ${fromState} -> ${toState} transition`);
+    };
+
     const transition = new Transition({
       debug: this.debug,
       listeners,
       check,
+      onError: onTransitionError,
       onMatch: transitioningOnMatch,
       onMismatch,
     });
@@ -461,9 +476,9 @@ export class StateMachine {
   async stopMachine() {
     this.debug && console.log('Stopping state machine...');
 
+    this.status = 'stopped';
     await this.exitCurrentState();
     await this.onStopCallback?.();
-    this.status = 'stopped';
 
     this.debug && console.log('State machine stopped');
   }
@@ -504,10 +519,6 @@ export class StateMachine {
    * Stops listening on the current state's transitions and exits the current state.
    */
   private async exitCurrentState() {
-    if (!this.isRunning) {
-      return;
-    }
-
     this.debug && console.log('exitCurrentState', this.currentState?.key);
 
     const currentTransitions =
@@ -547,7 +558,13 @@ export class StateMachine {
     const nextState = this.states.get(stateKey);
 
     if (!nextState) {
-      throw new Error(`State ${stateKey} not found`);
+      throw new UnknownError(
+        {
+          currentState: this.currentState,
+          nextState: stateKey,
+        },
+        `Machine next state not found`
+      );
     }
     if (this.currentState === nextState) {
       console.warn(`State ${stateKey} is already active. Skipping transition.`);
@@ -560,8 +577,33 @@ export class StateMachine {
       this.isRunning && (await this.enterState(stateKey));
     } catch (e) {
       this.currentState = undefined;
-      console.error(e);
-      throw new Error(`Could not enter state ${stateKey}`);
+      this.handleError(e, `Could not enter state ${stateKey}`);
+    }
+  }
+
+  private handleError(error: unknown, context: string) {
+    // Try to halt machine if it is still running
+    if (this.isRunning) {
+      const publicError = new AutomationError(
+        {
+          info: {
+            stateMachineId: this.id,
+            status: this.status,
+            currentState: this.currentState,
+          },
+          cause: error,
+        },
+        context ?? 'Error running state machine'
+      );
+      if (this.onError) {
+        this.onError(publicError);
+      } else {
+        // This throw will likely crash the server
+        throw publicError;
+      }
+
+      // Throwing when stopping could hide above error
+      this.stopMachine().catch(console.error);
     }
   }
 

packages/automation/src/lib/transitions/transition.ts

@@ -1,4 +1,5 @@
 import { Listener } from '../listeners';
+import { onError } from '../types';
 
 export type CheckFn = (values: (unknown | undefined)[]) => Promise<boolean>;
 export type resultFn = (values: (unknown | undefined)[]) => Promise<void>;
@@ -13,6 +14,7 @@ export interface BaseTransitionParams {
   check?: CheckFn;
   onMatch: resultFn;
   onMismatch?: resultFn;
+  onError?: onError;
 }
 
 export class Transition {
@@ -22,6 +24,7 @@ export class Transition {
   private readonly check?: CheckFn;
   private readonly onMatch: resultFn;
   private readonly onMismatch?: resultFn;
+  private readonly onError?: onError;
   private readonly queue: Values[] = [];
   private isProcessingQueue = false;
 
@@ -36,12 +39,14 @@ export class Transition {
     check,
     onMatch,
     onMismatch,
+    onError,
   }: BaseTransitionParams) {
     this.debug = debug ?? false;
     this.listeners = listeners;
     this.check = check;
     this.onMatch = onMatch;
     this.onMismatch = onMismatch;
+    this.onError = onError;
     this.values = new Array(listeners.length).fill(undefined);
     this.setupListeners();
   }
@@ -60,59 +65,84 @@ export class Transition {
         // Process the queue
         await this.processQueue();
       });
+      listener.onError?.(this.onError);
     });
   }
 
   /**
    * Starts all listeners for this transition.
    */
   async startListening() {
-    this.debug && console.log('startListening');
-    await Promise.all(this.listeners.map((listener) => listener.start()));
-
-    if (!this.listeners.length) {
-      // If the transition does not have any listeners it will never emit. Therefore, we "match" automatically on next event loop
-      setTimeout(() => {
-        this.debug && console.log('Transition without listeners: auto match');
-        this.onMatch([]);
-      }, 0);
+    try {
+      this.debug && console.log('startListening');
+      await Promise.all(this.listeners.map((listener) => listener.start()));
+
+      if (!this.listeners.length) {
+        // If the transition does not have any listeners it will never emit. Therefore, we "match" automatically on next event loop
+        setTimeout(() => {
+          this.debug && console.log('Transition without listeners: auto match');
+          this.onMatch([]);
+        }, 0);
+      }
+    } catch (e) {
+      if (this.onError) {
+        this.onError(e);
+      } else {
+        throw e;
+      }
     }
   }
 
   /**
    * Stops all listeners for this transition.
    */
   async stopListening() {
-    this.debug && console.log('stopListening');
-    this.queue.length = 0; // Flush the queue as there might be more value arrays to check
-    await Promise.all(this.listeners.map((listener) => listener.stop()));
+    try {
+      this.debug && console.log('stopListening');
+      this.queue.length = 0; // Flush the queue as there might be more value arrays to check
+      await Promise.all(this.listeners.map((listener) => listener.stop()));
+    } catch (e) {
+      if (this.onError) {
+        this.onError(e);
+      } else {
+        throw e;
+      }
+    }
   }
 
   private async processQueue() {
-    // Prevent concurrent queue processing
-    if (this.isProcessingQueue) {
-      return;
-    }
-    this.isProcessingQueue = true;
+    try {
+      // Prevent concurrent queue processing
+      if (this.isProcessingQueue) {
+        return;
+      }
+      this.isProcessingQueue = true;
 
-    while (this.queue.length > 0) {
-      const currentValues = this.queue.shift();
+      while (this.queue.length > 0) {
+        const currentValues = this.queue.shift();
 
-      if (!currentValues) {
-        continue;
-      }
+        if (!currentValues) {
+          continue;
+        }
 
-      const isMatch = this.check ? await this.check(currentValues) : true;
+        const isMatch = this.check ? await this.check(currentValues) : true;
 
-      if (isMatch) {
-        this.debug && console.log('match', currentValues);
-        await this.onMatch?.(currentValues);
+        if (isMatch) {
+          this.debug && console.log('match', currentValues);
+          await this.onMatch?.(currentValues);
+        } else {
+          this.debug && console.log('mismatch', currentValues);
+          await this.onMismatch?.(currentValues);
+        }
+      }
+
+      this.isProcessingQueue = false; // Allow new queue processing
+    } catch (e) {
+      if (this.onError) {
+        this.onError(e);
       } else {
-        this.debug && console.log('mismatch', currentValues);
-        await this.onMismatch?.(currentValues);
+        throw e;
       }
     }
-
-    this.isProcessingQueue = false; // Allow new queue processing
   }
 }

packages/automation/src/lib/types.ts

@@ -6,6 +6,7 @@ import { LitNodeClient } from '@lit-protocol/lit-node-client';
 import { BaseTransitionParams } from './transitions';
 
 export type Address = `0x${string}`;
+export type onError = (error: unknown) => void;
 
 export type PKPInfo = {
   tokenId: string;
@@ -125,6 +126,7 @@ export interface BaseStateMachineParams {
   debug?: boolean;
   litContracts: LitContracts;
   litNodeClient: LitNodeClient;
+  onError?: (error: unknown, context?: string) => void;
   privateKey?: string;
   pkp?: PKPInfo;
 }

packages/constants/src/lib/errors.ts

@@ -217,6 +217,11 @@ export const LIT_ERROR: Record<string, ErrorConfig> = {
     code: 'transaction_error',
     kind: LitErrorKind.Unexpected,
   },
+  AUTOMATION_ERROR: {
+    name: 'AutomationError',
+    code: 'automation_error',
+    kind: LitErrorKind.Unexpected,
+  },
 };
 
 export const LIT_ERROR_CODE = {
@@ -292,6 +297,7 @@ const MultiError = VError.MultiError;
 export { MultiError };
 
 export const {
+  AutomationError,
   InitError,
   InvalidAccessControlConditions,
   InvalidArgumentException,