@aztec/archiver 0.0.1-commit.381b1a9 → 0.0.1-commit.3a4ae741b

package/src/store/contract_class_store.ts

@@ -2,14 +2,7 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
 import { BufferReader, numToUInt8, serializeToBuffer } from '@aztec/foundation/serialize';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
-import { FunctionSelector } from '@aztec/stdlib/abi';
-import type {
-  ContractClassPublic,
-  ContractClassPublicWithBlockNumber,
-  ExecutablePrivateFunctionWithMembershipProof,
-  UtilityFunctionWithMembershipProof,
-} from '@aztec/stdlib/contract';
-import { Vector } from '@aztec/stdlib/types';
+import type { ContractClassPublic, ContractClassPublicWithBlockNumber } from '@aztec/stdlib/contract';
 
 /**
  * LMDB-based contract class storage for the archiver.
@@ -29,11 +22,15 @@ export class ContractClassStore {
     blockNumber: number,
   ): Promise<void> {
     await this.db.transactionAsync(async () => {
-      await this.#contractClasses.setIfNotExists(
-        contractClass.id.toString(),
+      const key = contractClass.id.toString();
+      if (await this.#contractClasses.hasAsync(key)) {
+        throw new Error(`Contract class ${key} already exists, cannot add again at block ${blockNumber}`);
+      }
+      await this.#contractClasses.set(
+        key,
         serializeContractClassPublic({ ...contractClass, l2BlockNumber: blockNumber }),
       );
-      await this.#bytecodeCommitments.setIfNotExists(contractClass.id.toString(), bytecodeCommitment.toBuffer());
+      await this.#bytecodeCommitments.set(key, bytecodeCommitment.toBuffer());
     });
   }
 
@@ -60,37 +57,6 @@ export class ContractClassStore {
   async getContractClassIds(): Promise<Fr[]> {
     return (await toArray(this.#contractClasses.keysAsync())).map(key => Fr.fromHexString(key));
   }
-
-  async addFunctions(
-    contractClassId: Fr,
-    newPrivateFunctions: ExecutablePrivateFunctionWithMembershipProof[],
-    newUtilityFunctions: UtilityFunctionWithMembershipProof[],
-  ): Promise<boolean> {
-    await this.db.transactionAsync(async () => {
-      const existingClassBuffer = await this.#contractClasses.getAsync(contractClassId.toString());
-      if (!existingClassBuffer) {
-        throw new Error(`Unknown contract class ${contractClassId} when adding private functions to store`);
-      }
-
-      const existingClass = deserializeContractClassPublic(existingClassBuffer);
-      const { privateFunctions: existingPrivateFns, utilityFunctions: existingUtilityFns } = existingClass;
-
-      const updatedClass: Omit<ContractClassPublicWithBlockNumber, 'id'> = {
-        ...existingClass,
-        privateFunctions: [
-          ...existingPrivateFns,
-          ...newPrivateFunctions.filter(newFn => !existingPrivateFns.some(f => f.selector.equals(newFn.selector))),
-        ],
-        utilityFunctions: [
-          ...existingUtilityFns,
-          ...newUtilityFunctions.filter(newFn => !existingUtilityFns.some(f => f.selector.equals(newFn.selector))),
-        ],
-      };
-      await this.#contractClasses.set(contractClassId.toString(), serializeContractClassPublic(updatedClass));
-    });
-
-    return true;
-  }
 }
 
 function serializeContractClassPublic(contractClass: Omit<ContractClassPublicWithBlockNumber, 'id'>): Buffer {
@@ -98,83 +64,19 @@ function serializeContractClassPublic(contractClass: Omit<ContractClassPublicWit
     contractClass.l2BlockNumber,
     numToUInt8(contractClass.version),
     contractClass.artifactHash,
-    contractClass.privateFunctions.length,
-    contractClass.privateFunctions.map(serializePrivateFunction),
-    contractClass.utilityFunctions.length,
-    contractClass.utilityFunctions.map(serializeUtilityFunction),
     contractClass.packedBytecode.length,
     contractClass.packedBytecode,
     contractClass.privateFunctionsRoot,
   );
 }
 
-function serializePrivateFunction(fn: ExecutablePrivateFunctionWithMembershipProof): Buffer {
-  return serializeToBuffer(
-    fn.selector,
-    fn.vkHash,
-    fn.bytecode.length,
-    fn.bytecode,
-    fn.functionMetadataHash,
-    fn.artifactMetadataHash,
-    fn.utilityFunctionsTreeRoot,
-    new Vector(fn.privateFunctionTreeSiblingPath),
-    fn.privateFunctionTreeLeafIndex,
-    new Vector(fn.artifactTreeSiblingPath),
-    fn.artifactTreeLeafIndex,
-  );
-}
-
-function serializeUtilityFunction(fn: UtilityFunctionWithMembershipProof): Buffer {
-  return serializeToBuffer(
-    fn.selector,
-    fn.bytecode.length,
-    fn.bytecode,
-    fn.functionMetadataHash,
-    fn.artifactMetadataHash,
-    fn.privateFunctionsArtifactTreeRoot,
-    new Vector(fn.artifactTreeSiblingPath),
-    fn.artifactTreeLeafIndex,
-  );
-}
-
 function deserializeContractClassPublic(buffer: Buffer): Omit<ContractClassPublicWithBlockNumber, 'id'> {
   const reader = BufferReader.asReader(buffer);
   return {
     l2BlockNumber: reader.readNumber(),
     version: reader.readUInt8() as 1,
     artifactHash: reader.readObject(Fr),
-    privateFunctions: reader.readVector({ fromBuffer: deserializePrivateFunction }),
-    utilityFunctions: reader.readVector({ fromBuffer: deserializeUtilityFunction }),
     packedBytecode: reader.readBuffer(),
     privateFunctionsRoot: reader.readObject(Fr),
   };
 }
-
-function deserializePrivateFunction(buffer: Buffer | BufferReader): ExecutablePrivateFunctionWithMembershipProof {
-  const reader = BufferReader.asReader(buffer);
-  return {
-    selector: reader.readObject(FunctionSelector),
-    vkHash: reader.readObject(Fr),
-    bytecode: reader.readBuffer(),
-    functionMetadataHash: reader.readObject(Fr),
-    artifactMetadataHash: reader.readObject(Fr),
-    utilityFunctionsTreeRoot: reader.readObject(Fr),
-    privateFunctionTreeSiblingPath: reader.readVector(Fr),
-    privateFunctionTreeLeafIndex: reader.readNumber(),
-    artifactTreeSiblingPath: reader.readVector(Fr),
-    artifactTreeLeafIndex: reader.readNumber(),
-  };
-}
-
-function deserializeUtilityFunction(buffer: Buffer | BufferReader): UtilityFunctionWithMembershipProof {
-  const reader = BufferReader.asReader(buffer);
-  return {
-    selector: reader.readObject(FunctionSelector),
-    bytecode: reader.readBuffer(),
-    functionMetadataHash: reader.readObject(Fr),
-    artifactMetadataHash: reader.readObject(Fr),
-    privateFunctionsArtifactTreeRoot: reader.readObject(Fr),
-    artifactTreeSiblingPath: reader.readVector(Fr),
-    artifactTreeLeafIndex: reader.readNumber(),
-  };
-}

package/src/store/contract_instance_store.ts

@@ -27,11 +27,14 @@ export class ContractInstanceStore {
 
   addContractInstance(contractInstance: ContractInstanceWithAddress, blockNumber: number): Promise<void> {
     return this.db.transactionAsync(async () => {
-      await this.#contractInstances.set(
-        contractInstance.address.toString(),
-        new SerializableContractInstance(contractInstance).toBuffer(),
-      );
-      await this.#contractInstancePublishedAt.set(contractInstance.address.toString(), blockNumber);
+      const key = contractInstance.address.toString();
+      if (await this.#contractInstances.hasAsync(key)) {
+        throw new Error(
+          `Contract instance at ${key} already exists (deployed at block ${await this.#contractInstancePublishedAt.getAsync(key)}), cannot add again at block ${blockNumber}`,
+        );
+      }
+      await this.#contractInstances.set(key, new SerializableContractInstance(contractInstance).toBuffer());
+      await this.#contractInstancePublishedAt.set(key, blockNumber);
     });
   }
 

package/src/store/kv_archiver_store.ts

@@ -16,11 +16,10 @@ import {
 import type { CheckpointData, PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
 import type {
   ContractClassPublic,
+  ContractClassPublicWithCommitment,
   ContractDataSource,
   ContractInstanceUpdateWithAddress,
   ContractInstanceWithAddress,
-  ExecutablePrivateFunctionWithMembershipProof,
-  UtilityFunctionWithMembershipProof,
 } from '@aztec/stdlib/contract';
 import type { L1RollupConstants } from '@aztec/stdlib/epoch-helpers';
 import type { GetContractClassLogsResponse, GetPublicLogsResponse } from '@aztec/stdlib/interfaces/client';
@@ -37,7 +36,7 @@ import { ContractInstanceStore } from './contract_instance_store.js';
 import { LogStore } from './log_store.js';
 import { MessageStore } from './message_store.js';
 
-export const ARCHIVER_DB_VERSION = 5;
+export const ARCHIVER_DB_VERSION = 6;
 export const MAX_FUNCTION_SIGNATURES = 1000;
 export const MAX_FUNCTION_NAME_LEN = 256;
 
@@ -167,19 +166,14 @@ export class KVArchiverDataStore implements ContractDataSource {
 
   /**
    * Add new contract classes from an L2 block to the store's list.
-   * @param data - List of contract classes to be added.
-   * @param bytecodeCommitments - Bytecode commitments for the contract classes.
+   * @param data - List of contract classes (with bytecode commitments) to be added.
    * @param blockNumber - Number of the L2 block the contracts were registered in.
    * @returns True if the operation is successful.
    */
-  async addContractClasses(
-    data: ContractClassPublic[],
-    bytecodeCommitments: Fr[],
-    blockNumber: BlockNumber,
-  ): Promise<boolean> {
+  async addContractClasses(data: ContractClassPublicWithCommitment[], blockNumber: BlockNumber): Promise<boolean> {
     return (
       await Promise.all(
-        data.map((c, i) => this.#contractClassStore.addContractClass(c, bytecodeCommitments[i], blockNumber)),
+        data.map(c => this.#contractClassStore.addContractClass(c, c.publicBytecodeCommitment, blockNumber)),
       )
     ).every(Boolean);
   }
@@ -194,15 +188,6 @@ export class KVArchiverDataStore implements ContractDataSource {
     return this.#contractClassStore.getBytecodeCommitment(contractClassId);
   }
 
-  /** Adds private functions to a contract class. */
-  addFunctions(
-    contractClassId: Fr,
-    privateFunctions: ExecutablePrivateFunctionWithMembershipProof[],
-    utilityFunctions: UtilityFunctionWithMembershipProof[],
-  ): Promise<boolean> {
-    return this.#contractClassStore.addFunctions(contractClassId, privateFunctions, utilityFunctions);
-  }
-
   /**
    * Add new contract instances from an L2 block to the store's list.
    * @param data - List of contract instances to be added.
@@ -245,14 +230,14 @@ export class KVArchiverDataStore implements ContractDataSource {
   }
 
   /**
-   * Append new proposed blocks to the store's list.
-   * These are uncheckpointed blocks that have been proposed by the sequencer but not yet included in a checkpoint on L1.
+   * Append a new proposed block to the store.
+   * This is an uncheckpointed block that has been proposed by the sequencer but not yet included in a checkpoint on L1.
    * For checkpointed blocks (already published to L1), use addCheckpoints() instead.
-   * @param blocks - The proposed L2 blocks to be added to the store.
+   * @param block - The proposed L2 block to be added to the store.
    * @returns True if the operation is successful.
    */
-  addProposedBlocks(blocks: L2Block[], opts: { force?: boolean; checkpointNumber?: number } = {}): Promise<boolean> {
-    return this.#blockStore.addProposedBlocks(blocks, opts);
+  addProposedBlock(block: L2Block, opts: { force?: boolean } = {}): Promise<boolean> {
+    return this.#blockStore.addProposedBlock(block, opts);
   }
 
   /**
@@ -474,10 +459,11 @@ export class KVArchiverDataStore implements ContractDataSource {
    * array implies no logs match that tag.
    * @param tags - The tags to search for.
    * @param page - The page number (0-indexed) for pagination. Returns at most 10 logs per tag per page.
+   * @param upToBlockNumber - If set, only return logs from blocks up to and including this block number.
    */
-  getPrivateLogsByTags(tags: SiloedTag[], page?: number): Promise<TxScopedL2Log[][]> {
+  getPrivateLogsByTags(tags: SiloedTag[], page?: number, upToBlockNumber?: BlockNumber): Promise<TxScopedL2Log[][]> {
     try {
-      return this.#logStore.getPrivateLogsByTags(tags, page);
+      return this.#logStore.getPrivateLogsByTags(tags, page, upToBlockNumber);
     } catch (err) {
       return Promise.reject(err);
     }
@@ -489,14 +475,16 @@ export class KVArchiverDataStore implements ContractDataSource {
    * @param contractAddress - The contract address to search logs for.
    * @param tags - The tags to search for.
    * @param page - The page number (0-indexed) for pagination. Returns at most 10 logs per tag per page.
+   * @param upToBlockNumber - If set, only return logs from blocks up to and including this block number.
    */
   getPublicLogsByTagsFromContract(
     contractAddress: AztecAddress,
     tags: Tag[],
     page?: number,
+    upToBlockNumber?: BlockNumber,
   ): Promise<TxScopedL2Log[][]> {
     try {
-      return this.#logStore.getPublicLogsByTagsFromContract(contractAddress, tags, page);
+      return this.#logStore.getPublicLogsByTagsFromContract(contractAddress, tags, page, upToBlockNumber);
     } catch (err) {
       return Promise.reject(err);
     }
@@ -603,6 +591,11 @@ export class KVArchiverDataStore implements ContractDataSource {
     return this.#messageStore.rollbackL1ToL2MessagesToCheckpoint(targetCheckpointNumber);
   }
 
+  /** Persists the inbox tree-in-progress checkpoint number from L1 state. */
+  public setInboxTreeInProgress(value: bigint): Promise<void> {
+    return this.#messageStore.setInboxTreeInProgress(value);
+  }
+
   /** Returns an async iterator to all L1 to L2 messages on the range. */
   public iterateL1ToL2Messages(range: CustomRange<bigint> = {}): AsyncIterableIterator<InboxMessage> {
     return this.#messageStore.iterateL1ToL2Messages(range);

package/src/store/log_store.ts

@@ -22,6 +22,7 @@ import {
 } from '@aztec/stdlib/logs';
 import { TxHash } from '@aztec/stdlib/tx';
 
+import { OutOfOrderLogInsertionError } from '../errors.js';
 import type { BlockStore } from './block_store.js';
 
 /**
@@ -165,10 +166,21 @@ export class LogStore {
 
     for (const taggedLogBuffer of currentPrivateTaggedLogs) {
       if (taggedLogBuffer.logBuffers && taggedLogBuffer.logBuffers.length > 0) {
-        privateTaggedLogs.set(
-          taggedLogBuffer.tag,
-          taggedLogBuffer.logBuffers!.concat(privateTaggedLogs.get(taggedLogBuffer.tag)!),
-        );
+        const newLogs = privateTaggedLogs.get(taggedLogBuffer.tag)!;
+        if (newLogs.length === 0) {
+          continue;
+        }
+        const lastExisting = TxScopedL2Log.fromBuffer(taggedLogBuffer.logBuffers.at(-1)!);
+        const firstNew = TxScopedL2Log.fromBuffer(newLogs[0]);
+        if (lastExisting.blockNumber > firstNew.blockNumber) {
+          throw new OutOfOrderLogInsertionError(
+            'private',
+            taggedLogBuffer.tag,
+            lastExisting.blockNumber,
+            firstNew.blockNumber,
+          );
+        }
+        privateTaggedLogs.set(taggedLogBuffer.tag, taggedLogBuffer.logBuffers.concat(newLogs));
       }
     }
 
@@ -200,10 +212,21 @@ export class LogStore {
 
     for (const taggedLogBuffer of currentPublicTaggedLogs) {
       if (taggedLogBuffer.logBuffers && taggedLogBuffer.logBuffers.length > 0) {
-        publicTaggedLogs.set(
-          taggedLogBuffer.tag,
-          taggedLogBuffer.logBuffers!.concat(publicTaggedLogs.get(taggedLogBuffer.tag)!),
-        );
+        const newLogs = publicTaggedLogs.get(taggedLogBuffer.tag)!;
+        if (newLogs.length === 0) {
+          continue;
+        }
+        const lastExisting = TxScopedL2Log.fromBuffer(taggedLogBuffer.logBuffers.at(-1)!);
+        const firstNew = TxScopedL2Log.fromBuffer(newLogs[0]);
+        if (lastExisting.blockNumber > firstNew.blockNumber) {
+          throw new OutOfOrderLogInsertionError(
+            'public',
+            taggedLogBuffer.tag,
+            lastExisting.blockNumber,
+            firstNew.blockNumber,
+          );
+        }
+        publicTaggedLogs.set(taggedLogBuffer.tag, taggedLogBuffer.logBuffers.concat(newLogs));
       }
     }
 
@@ -353,17 +376,30 @@ export class LogStore {
    * array implies no logs match that tag.
    * @param tags - The tags to search for.
    * @param page - The page number (0-indexed) for pagination.
+   * @param upToBlockNumber - If set, only return logs from blocks up to and including this block number.
    * @returns An array of log arrays, one per tag. Returns at most MAX_LOGS_PER_TAG logs per tag per page. If
    * MAX_LOGS_PER_TAG logs are returned for a tag, the caller should fetch the next page to check for more logs.
    */
-  async getPrivateLogsByTags(tags: SiloedTag[], page: number = 0): Promise<TxScopedL2Log[][]> {
+  async getPrivateLogsByTags(
+    tags: SiloedTag[],
+    page: number = 0,
+    upToBlockNumber?: BlockNumber,
+  ): Promise<TxScopedL2Log[][]> {
     const logs = await Promise.all(tags.map(tag => this.#privateLogsByTag.getAsync(tag.toString())));
+
     const start = page * MAX_LOGS_PER_TAG;
     const end = start + MAX_LOGS_PER_TAG;
 
-    return logs.map(
-      logBuffers => logBuffers?.slice(start, end).map(logBuffer => TxScopedL2Log.fromBuffer(logBuffer)) ?? [],
-    );
+    return logs.map(logBuffers => {
+      const deserialized = logBuffers?.slice(start, end).map(buf => TxScopedL2Log.fromBuffer(buf)) ?? [];
+      if (upToBlockNumber !== undefined) {
+        const cutoff = deserialized.findIndex(log => log.blockNumber > upToBlockNumber);
+        if (cutoff !== -1) {
+          return deserialized.slice(0, cutoff);
+        }
+      }
+      return deserialized;
+    });
   }
 
   /**
@@ -372,6 +408,7 @@ export class LogStore {
    * @param contractAddress - The contract address to search logs for.
    * @param tags - The tags to search for.
    * @param page - The page number (0-indexed) for pagination.
+   * @param upToBlockNumber - If set, only return logs from blocks up to and including this block number.
    * @returns An array of log arrays, one per tag. Returns at most MAX_LOGS_PER_TAG logs per tag per page. If
    * MAX_LOGS_PER_TAG logs are returned for a tag, the caller should fetch the next page to check for more logs.
    */
@@ -379,6 +416,7 @@ export class LogStore {
     contractAddress: AztecAddress,
     tags: Tag[],
     page: number = 0,
+    upToBlockNumber?: BlockNumber,
   ): Promise<TxScopedL2Log[][]> {
     const logs = await Promise.all(
       tags.map(tag => {
@@ -389,9 +427,16 @@ export class LogStore {
     const start = page * MAX_LOGS_PER_TAG;
     const end = start + MAX_LOGS_PER_TAG;
 
-    return logs.map(
-      logBuffers => logBuffers?.slice(start, end).map(logBuffer => TxScopedL2Log.fromBuffer(logBuffer)) ?? [],
-    );
+    return logs.map(logBuffers => {
+      const deserialized = logBuffers?.slice(start, end).map(buf => TxScopedL2Log.fromBuffer(buf)) ?? [];
+      if (upToBlockNumber !== undefined) {
+        const cutoff = deserialized.findIndex(log => log.blockNumber > upToBlockNumber);
+        if (cutoff !== -1) {
+          return deserialized.slice(0, cutoff);
+        }
+      }
+      return deserialized;
+    });
   }
 
   /**

package/src/store/message_store.ts

@@ -14,6 +14,7 @@ import {
 } from '@aztec/kv-store';
 import { InboxLeaf } from '@aztec/stdlib/messaging';
 
+import { L1ToL2MessagesNotReadyError } from '../errors.js';
 import {
   type InboxMessage,
   deserializeInboxMessage,
@@ -40,6 +41,8 @@ export class MessageStore {
   #lastSynchedL1Block: AztecAsyncSingleton<Buffer>;
   /** Stores total messages stored */
   #totalMessageCount: AztecAsyncSingleton<bigint>;
+  /** Stores the checkpoint number whose message tree is currently being filled on L1. */
+  #inboxTreeInProgress: AztecAsyncSingleton<bigint>;
 
   #log = createLogger('archiver:message_store');
 
@@ -48,6 +51,7 @@ export class MessageStore {
     this.#l1ToL2MessageIndices = db.openMap('archiver_l1_to_l2_message_indices');
     this.#lastSynchedL1Block = db.openSingleton('archiver_last_l1_block_id');
     this.#totalMessageCount = db.openSingleton('archiver_l1_to_l2_message_count');
+    this.#inboxTreeInProgress = db.openSingleton('archiver_inbox_tree_in_progress');
   }
 
   public async getTotalL1ToL2MessageCount(): Promise<bigint> {
@@ -185,7 +189,22 @@ export class MessageStore {
     return msg ? deserializeInboxMessage(msg) : undefined;
   }
 
+  /** Returns the inbox tree-in-progress checkpoint number from L1, or undefined if not yet set. */
+  public getInboxTreeInProgress(): Promise<bigint | undefined> {
+    return this.#inboxTreeInProgress.getAsync();
+  }
+
+  /** Persists the inbox tree-in-progress checkpoint number from L1 state. */
+  public async setInboxTreeInProgress(value: bigint): Promise<void> {
+    await this.#inboxTreeInProgress.set(value);
+  }
+
   public async getL1ToL2Messages(checkpointNumber: CheckpointNumber): Promise<Fr[]> {
+    const treeInProgress = await this.#inboxTreeInProgress.getAsync();
+    if (treeInProgress !== undefined && BigInt(checkpointNumber) >= treeInProgress) {
+      throw new L1ToL2MessagesNotReadyError(checkpointNumber, treeInProgress);
+    }
+
     const messages: Fr[] = [];
 
     const [startIndex, endIndex] = InboxLeaf.indexRangeForCheckpoint(checkpointNumber);

package/src/test/fake_l1_state.ts

@@ -1,5 +1,6 @@
 import type { BlobClientInterface } from '@aztec/blob-client/client';
 import { type Blob, getBlobsPerL1Block, getPrefixedEthBlobCommitments } from '@aztec/blob-lib';
+import { INITIAL_CHECKPOINT_NUMBER } from '@aztec/constants';
 import type { CheckpointProposedLog, InboxContract, MessageSentLog, RollupContract } from '@aztec/ethereum/contracts';
 import { MULTI_CALL_3_ADDRESS } from '@aztec/ethereum/contracts';
 import type { ViemPublicClient } from '@aztec/ethereum/types';
@@ -450,13 +451,22 @@ export class FakeL1State {
   createMockInboxContract(_publicClient: MockProxy<ViemPublicClient>): MockProxy<InboxContract> {
     const mockInbox = mock<InboxContract>();
 
-    mockInbox.getState.mockImplementation(() =>
-      Promise.resolve({
+    mockInbox.getState.mockImplementation(() => {
+      // treeInProgress must be > any sealed checkpoint. On L1, a checkpoint can only be proposed
+      // after its messages are sealed, so treeInProgress > checkpointNumber for all published checkpoints.
+      const maxFromMessages =
+        this.messages.length > 0 ? Math.max(...this.messages.map(m => Number(m.checkpointNumber))) + 1 : 0;
+      const maxFromCheckpoints =
+        this.checkpoints.length > 0
+          ? Math.max(...this.checkpoints.filter(cp => !cp.pruned).map(cp => Number(cp.checkpointNumber))) + 1
+          : 0;
+      const treeInProgress = Math.max(maxFromMessages, maxFromCheckpoints, INITIAL_CHECKPOINT_NUMBER);
+      return Promise.resolve({
         messagesRollingHash: this.messagesRollingHash,
         totalMessagesInserted: BigInt(this.messages.length),
-        treeInProgress: 0n,
-      }),
-    );
+        treeInProgress: BigInt(treeInProgress),
+      });
+    });
 
     // Mock the wrapper methods for fetching message events
     mockInbox.getMessageSentEvents.mockImplementation((fromBlock: bigint, toBlock: bigint) =>