Mitigate issues with missing awaits in state method calls

packages/common/src/index.ts

@@ -17,3 +17,4 @@ export * from "./trees/InMemoryMerkleTreeStorage";
 export * from "./trees/RollupMerkleTree";
 export * from "./events/EventEmitterProxy";
 export * from "./trees/MockAsyncMerkleStore";
+export * from "./operationQueue";

packages/common/src/operationQueue.ts

@@ -0,0 +1,37 @@
+export type AsyncOperation<T> = () => Promise<T>;
+
+export class OperationQueue {
+  private queue: Promise<void>;
+
+  constructor() {
+    // Start with a resolved promise
+    this.queue = Promise.resolve();
+  }
+
+  /**
+   * Queue an operation to be executed
+   * @param operation - The operation to queue
+   * @returns A promise that resolves when the operation is completed
+   */
+  public queueOperation<T>(operation: AsyncOperation<T>): Promise<T> {
+    // Chain the operation to the end of the queue and store the result
+    // of the queued operation to return to the caller if he wants to wait
+    // for it
+    const result = this.queue.then(() => operation());
+
+    // Update the queue and discard any errors
+    this.queue = result.then(
+      () => undefined,
+      () => undefined
+    );
+    return result;
+  }
+
+  /**
+   * Wait for all operations to complete
+   * @returns A promise that resolves when all operations are completed
+   */
+  public async onCompleted(): Promise<void> {
+    await this.queue;
+  }
+}

packages/module/src/method/runtimeMethod.ts

@@ -89,6 +89,10 @@ export function toWrappedMethod(
     ...args
   ): Promise<MethodPublicOutput> => {
     await Reflect.apply(moduleMethod, this, args);
+    // await pending state operations to complete
+    await container
+      .resolve(RuntimeMethodExecutionContext)
+      .operationQueue.onCompleted();
     const {
       result: { stateTransitions, status, events },
     } = executionContext.current();
@@ -285,6 +289,8 @@ function runtimeMethodInternal(options: {
       let result: unknown;
       try {
         result = await Reflect.apply(simulatedMethod, this, args);
+        // await pending state operations to complete
+        await executionContext.operationQueue.onCompleted();
       } finally {
         executionContext.afterMethod();
       }

packages/protocol/src/state/State.ts

@@ -130,18 +130,22 @@ export class State<Value> extends Mixin(WithPath, WithStateServiceProvider) {
    *
    * @returns Option representation of the current state.
    */
-  public async get(): Promise<Option<Value>> {
-    const option = await this.witnessFromState();
+  public async get() {
+    return await container
+      .resolve(RuntimeMethodExecutionContext)
+      .operationQueue.queueOperation(async () => {
+        const option = await this.witnessFromState();
 
-    this.hasPathOrFail();
+        this.hasPathOrFail();
 
-    const stateTransition = StateTransition.from(this.path, option);
+        const stateTransition = StateTransition.from(this.path, option);
 
-    container
-      .resolve(RuntimeMethodExecutionContext)
-      .addStateTransition(stateTransition);
+        container
+          .resolve(RuntimeMethodExecutionContext)
+          .addStateTransition(stateTransition);
 
-    return option;
+        return option;
+      });
   }
 
   /**
@@ -156,20 +160,24 @@ export class State<Value> extends Mixin(WithPath, WithStateServiceProvider) {
    * @param value - Value to be set as the current state
    */
   public async set(value: Value) {
-    // link the transition to the current state
-    const fromOption = await this.witnessFromState();
-    const toOption = Option.fromValue(value, this.valueType);
-
-    this.hasPathOrFail();
-
-    const stateTransition = StateTransition.fromTo(
-      this.path,
-      fromOption,
-      toOption
-    );
-
-    container
+    return await container
       .resolve(RuntimeMethodExecutionContext)
-      .addStateTransition(stateTransition);
+      .operationQueue.queueOperation(async () => {
+        // link the transition to the current state
+        const fromOption = await this.witnessFromState();
+        const toOption = Option.fromValue(value, this.valueType);
+
+        this.hasPathOrFail();
+
+        const stateTransition = StateTransition.fromTo(
+          this.path,
+          fromOption,
+          toOption
+        );
+
+        container
+          .resolve(RuntimeMethodExecutionContext)
+          .addStateTransition(stateTransition);
+      });
   }
 }

packages/protocol/src/state/context/RuntimeMethodExecutionContext.ts

@@ -1,6 +1,7 @@
 import { Bool, FlexibleProvablePure, Provable, Struct } from "o1js";
 import { singleton } from "tsyringe";
 import {
+  OperationQueue,
   ProvableMethodExecutionContext,
   ProvableMethodExecutionResult,
 } from "@proto-kit/common";
@@ -63,6 +64,8 @@ export class RuntimeMethodExecutionContext extends ProvableMethodExecutionContex
 
   private isSimulated: boolean = false;
 
+  public operationQueue = new OperationQueue();
+
   private assertSetupCalled(): asserts this is {
     input: RuntimeMethodExecutionData;
   } {