Solution to Failed TX Hooks

packages/library/src/hooks/TransactionFeeHook.ts

@@ -9,12 +9,13 @@ import {
   BeforeTransactionHookArguments,
   ProvableTransactionHook,
   PublicKeyOption,
+  StateMap,
 } from "@proto-kit/protocol";
 import { Field, Provable, PublicKey } from "o1js";
 import { noop } from "@proto-kit/common";
 
 import { UInt64 } from "../math/UInt64";
-import { Balance, TokenId } from "../runtime/Balances";
+import { Balance, BalancesKey, TokenId } from "../runtime/Balances";
 
 import {
   MethodFeeConfigData,
@@ -29,6 +30,8 @@ interface Balances {
     to: PublicKey,
     amount: Balance
   ) => Promise<void>;
+
+  balances: StateMap<BalancesKey, Balance>;
 }
 
 export interface TransactionFeeHookConfig
@@ -159,4 +162,24 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
   public async afterTransaction(): Promise<void> {
     noop();
   }
+
+  public async removeTransactionWhen(
+    args: BeforeTransactionHookArguments
+  ): Promise<boolean> {
+    const feeConfig = this.feeAnalyzer.getFeeConfig(
+      args.transaction.methodId.toBigInt()
+    );
+
+    const fee = this.getFee(feeConfig);
+
+    const tokenId = new TokenId(this.config.tokenId);
+    const feeRecipient = PublicKey.fromBase58(this.config.feeRecipient);
+
+    const balanceAvailable = await this.balances.balances.get({
+      tokenId,
+      address: feeRecipient,
+    });
+
+    return balanceAvailable.orElse(Balance.from(0)).lessThan(fee).toBoolean();
+  }
 }

packages/persistance/prisma/migrations/20251105204729_hooksstatus/migration.sql

@@ -0,0 +1,8 @@
+/*
+  Warnings:
+
+  - Added the required column `hooksStatus` to the `TransactionExecutionResult` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- AlterTable
+ALTER TABLE "TransactionExecutionResult" ADD COLUMN     "hooksStatus" BOOLEAN NOT NULL;

packages/persistance/prisma/schema.prisma

@@ -62,6 +62,7 @@ model TransactionExecutionResult {
   status           Boolean
   statusMessage    String?
   events           Json    @db.Json
+  hooksStatus      Boolean
 
   tx     Transaction @relation(fields: [txHash], references: [hash])
   txHash String      @id

packages/persistance/src/services/prisma/PrismaBlockStorage.ts

@@ -122,6 +122,7 @@ export class PrismaBlockStorage implements BlockQueue, BlockStorage {
             data: transactions.map((tx) => {
               return {
                 status: tx.status,
+                hooksStatus: tx.hooksStatus,
                 statusMessage: tx.statusMessage,
                 txHash: tx.txHash,
 

packages/persistance/src/services/prisma/PrismaTransactionStorage.ts

@@ -35,6 +35,22 @@ export class PrismaTransactionStorage implements TransactionStorage {
     return txs.map((tx) => this.transactionMapper.mapIn(tx));
   }
 
+  public async removeTx(hashes: string[], type: "included" | "dropped") {
+    // In our schema, included txs are simply just linked with blocks, so we only
+    // need to delete if we drop a tx
+    if (type === "dropped") {
+      const { prismaClient } = this.connection;
+
+      await prismaClient.transaction.deleteMany({
+        where: {
+          hash: {
+            in: hashes,
+          },
+        },
+      });
+    }
+  }
+
   public async pushUserTransaction(tx: PendingTransaction): Promise<boolean> {
     const { prismaClient } = this.connection;
 

packages/persistance/src/services/prisma/mappers/TransactionMapper.ts

@@ -66,6 +66,7 @@ export class TransactionExecutionResultMapper
     return {
       tx: this.transactionMapper.mapIn(input[1]),
       status: Bool(executionResult.status),
+      hooksStatus: Bool(executionResult.hooksStatus),
       statusMessage: executionResult.statusMessage ?? undefined,
       stateTransitions: this.stBatchMapper.mapIn(
         executionResult.stateTransitions
@@ -80,6 +81,7 @@ export class TransactionExecutionResultMapper
     const tx = this.transactionMapper.mapOut(input.tx);
     const executionResult = {
       status: input.status.toBoolean(),
+      hooksStatus: input.hooksStatus.toBoolean(),
       statusMessage: input.statusMessage ?? null,
       stateTransitions: this.stBatchMapper.mapOut(input.stateTransitions),
       events: this.eventArrayMapper.mapOut(input.events),

packages/protocol/src/hooks/AccountStateHook.ts

@@ -54,4 +54,19 @@ export class AccountStateHook extends ProvableTransactionHook {
   public async afterTransaction() {
     noop();
   }
+
+  // Under these conditions we want the tx removed from the mempool.
+  public async removeTransactionWhen({
+    transaction,
+  }: BeforeTransactionHookArguments): Promise<boolean> {
+    const sender = transaction.sender.value;
+
+    const aso = await this.accountState.get(sender);
+
+    const accountState = aso.orElse(new AccountState({ nonce: UInt64.zero }));
+
+    const currentNonce = accountState.nonce;
+
+    return transaction.nonce.value.lessThan(currentNonce).toBoolean();
+  }
 }

packages/protocol/src/protocol/ProvableTransactionHook.ts

@@ -74,4 +74,10 @@ export abstract class ProvableTransactionHook<
   public abstract afterTransaction(
     execution: AfterTransactionHookArguments
   ): Promise<void>;
+
+  public async removeTransactionWhen(
+    execution: BeforeTransactionHookArguments
+  ): Promise<boolean> {
+    return false;
+  }
 }

packages/protocol/src/prover/block/BlockProver.ts

@@ -171,7 +171,8 @@ export class BlockProverProgrammable extends ZkProgrammable<
     // Apply beforeTransaction hook state transitions
     const beforeBatch = await this.executeTransactionHooks(
       async (module, args) => await module.beforeTransaction(args),
-      beforeTxHookArguments
+      beforeTxHookArguments,
+      isMessage
     );
 
     state = this.addTransactionToBundle(
@@ -200,7 +201,8 @@ export class BlockProverProgrammable extends ZkProgrammable<
 
     const afterBatch = await this.executeTransactionHooks(
       async (module, args) => await module.afterTransaction(args),
-      afterTxHookArguments
+      afterTxHookArguments,
+      isMessage
     );
     state.pendingSTBatches.push(afterBatch);
 
@@ -266,20 +268,31 @@ export class BlockProverProgrammable extends ZkProgrammable<
     T extends BeforeTransactionHookArguments | AfterTransactionHookArguments,
   >(
     hook: (module: ProvableTransactionHook<unknown>, args: T) => Promise<void>,
-    hookArguments: T
+    hookArguments: T,
+    isMessage: Bool
   ) {
-    const { batch } = await this.executeHooks(hookArguments, async () => {
-      for (const module of this.transactionHooks) {
-        // eslint-disable-next-line no-await-in-loop
-        await hook(module, hookArguments);
-      }
-    });
+    const { batch, rawStatus } = await this.executeHooks(
+      hookArguments,
+      async () => {
+        for (const module of this.transactionHooks) {
+          // eslint-disable-next-line no-await-in-loop
+          await hook(module, hookArguments);
+        }
+      },
+      isMessage
+    );
+
+    // This is going to set applied to false in case the hook fails
+    // (that's only possible for messages though as others are hard-asserted)
+    batch.applied = rawStatus;
+
     return batch;
   }
 
   private async executeHooks<T>(
     contextArguments: RuntimeMethodExecutionData,
-    method: () => Promise<T>
+    method: () => Promise<T>,
+    isMessage: Bool | undefined = undefined
   ) {
     const executionContext = container.resolve(RuntimeMethodExecutionContext);
     executionContext.clear();
@@ -297,11 +310,23 @@ export class BlockProverProgrammable extends ZkProgrammable<
     const { stateTransitions, status, statusMessage } =
       executionContext.current().result;
 
-    status.assertTrue(`Transaction hook call failed: ${statusMessage ?? "-"}`);
+    // See https://github.com/proto-kit/framework/issues/321 for why we do this here
+    if (isMessage !== undefined) {
+      // isMessage is defined for all tx hooks
+      status
+        .or(isMessage)
+        .assertTrue(
+          `Transaction hook call failed for non-message tx: ${statusMessage ?? "-"}`
+        );
+    } else {
+      // isMessage is undefined for all block hooks
+      status.assertTrue(`Block hook call failed: ${statusMessage ?? "-"}`);
+    }
 
     return {
       batch: this.constructBatch(stateTransitions, Bool(true)),
       result,
+      rawStatus: status,
     };
   }
 

packages/protocol/src/state/assert/assert.ts

@@ -19,7 +19,7 @@ export function assert(condition: Bool, message?: string | (() => string)) {
   const status = condition.and(previousStatus);
 
   Provable.asProver(() => {
-    if (!condition.toBoolean()) {
+    if (!condition.toBoolean() && previousStatus.toBoolean()) {
       const messageString =
         message !== undefined && typeof message === "function"
           ? message()

packages/sdk/test/fees-failures.test.ts

@@ -60,6 +60,10 @@ describe("fee errors due to limited funds in sender accounts", () => {
     appChain.setSigner(senderKey);
   });
 
+  afterAll(async () => {
+    await appChain.close();
+  });
+
   it("should allow a free faucet transaction", async () => {
     expect.assertions(2);
 
@@ -105,8 +109,7 @@ describe("fee errors due to limited funds in sender accounts", () => {
     await appChain.produceBlock();
 
     expect(logSpy).toHaveBeenCalledWith(
-      "Error in inclusion of tx, skipping",
-      Error("Protocol hooks not executable: From balance is insufficient")
+      "Error in inclusion of tx, removing as to removeWhen hooks: Protocol hooks not executable: From balance is insufficient"
     );
 
     const balance = await appChain.query.runtime.Balances.balances.get(

packages/sdk/test/fees-multi-zkprograms.test.ts

@@ -184,12 +184,21 @@ describe("check fee analyzer", () => {
           },
         },
       },
+      Sequencer: {
+        Mempool: {
+          validationEnabled: true,
+        },
+      },
     });
 
     await appChain.start();
     appChain.setSigner(senderKey);
   });
 
+  afterAll(async () => {
+    await appChain.close();
+  });
+
   it("with multiple zk programs", async () => {
     expect.assertions(12);
     const testModule1 = appChain.runtime.resolve("TestModule1");

packages/sequencer/src/mempool/Mempool.ts

@@ -18,4 +18,6 @@ export interface Mempool<Events extends MempoolEvents = MempoolEvents>
    * Retrieve all transactions that are currently in the mempool
    */
   getTxs: (limit?: number) => Promise<PendingTransaction[]>;
+
+  removeTxs: (included: string[], dropped: string[]) => Promise<void>;
 }

packages/sequencer/src/mempool/private/PrivateMempool.ts

@@ -110,16 +110,22 @@ export class PrivateMempool
     return result?.result.afterNetworkState;
   }
 
+  public async removeTxs(included: string[], dropped: string[]) {
+    await this.transactionStorage.removeTx(included, "included");
+    await this.transactionStorage.removeTx(dropped, "dropped");
+  }
+
   @trace("mempool.get_txs")
   public async getTxs(limit?: number): Promise<PendingTransaction[]> {
+    // TODO Add limit to the storage (or do something smarter entirely)
     const txs = await this.transactionStorage.getPendingUserTransactions();
 
     const baseCachedStateService = new CachedStateService(this.stateService);
 
     const networkState =
       (await this.getStagedNetworkState()) ?? NetworkState.empty();
 
-    const validationEnabled = this.config.validationEnabled ?? true;
+    const validationEnabled = this.config.validationEnabled ?? false;
     const sortedTxs = validationEnabled
       ? await this.checkTxValid(
           txs,
@@ -128,7 +134,7 @@ export class PrivateMempool
           networkState,
           limit
         )
-      : txs;
+      : txs.slice(0, limit);
 
     this.protocol.stateServiceProvider.popCurrentStateService();
     return sortedTxs;
@@ -188,6 +194,7 @@ export class PrivateMempool
       executionContext.setup(contextInputs);
 
       const signedTransaction = tx.toProtocolTransaction();
+
       // eslint-disable-next-line no-await-in-loop
       await this.accountStateHook.beforeTransaction({
         networkState: networkState,
@@ -218,6 +225,26 @@ export class PrivateMempool
           queue = queue.filter(distinctByPredicate((a, b) => a === b));
         }
       } else {
+        // eslint-disable-next-line no-await-in-loop
+        const removeTxWhen = await this.accountStateHook.removeTransactionWhen({
+          networkState: networkState,
+          transaction: signedTransaction.transaction,
+          signature: signedTransaction.signature,
+          prover: proverState,
+        });
+        if (removeTxWhen) {
+          // eslint-disable-next-line no-await-in-loop
+          await this.transactionStorage.removeTx(
+            [tx.hash().toString()],
+            "dropped"
+          );
+          log.trace(
+            `Deleting tx ${tx.hash().toString()}  from mempool because removeTransactionWhen condition is satisfied`
+          );
+          // eslint-disable-next-line no-continue
+          continue;
+        }
+
         log.trace(
           `Skipped tx ${tx.hash().toString()} because ${statusMessage}`
         );

packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts

@@ -251,6 +251,16 @@ export class BlockProducerModule extends SequencerModule<BlockConfig> {
           height: block.height.toString(),
         }
       );
+
+      // Remove included or dropped txs, leave skipped ones alone
+      await this.mempool.removeTxs(
+        blockResult.includedTxs
+          .filter((x) => x.type === "included")
+          .map((x) => x.hash),
+        blockResult.includedTxs
+          .filter((x) => x.type === "shouldRemove")
+          .map((x) => x.hash)
+      );
     }
 
     this.productionInProgress = false;