@aztec/prover-client 0.0.1-commit.b655e406 → 0.0.1-commit.d1f2d6c

package/src/orchestrator/orchestrator.ts

@@ -1,4 +1,4 @@
-import { BatchedBlob, FinalBlobBatchingChallenges, SpongeBlob } from '@aztec/blob-lib';
+import { BatchedBlob, FinalBlobBatchingChallenges, SpongeBlob } from '@aztec/blob-lib/types';
 import {
   L1_TO_L2_MSG_SUBTREE_HEIGHT,
   L1_TO_L2_MSG_SUBTREE_ROOT_SIBLING_PATH_LENGTH,
@@ -6,22 +6,23 @@ import {
   NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
   NUM_BASE_PARITY_PER_ROOT_PARITY,
 } from '@aztec/constants';
+import { BlockNumber, EpochNumber } from '@aztec/foundation/branded-types';
 import { padArrayEnd } from '@aztec/foundation/collection';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { AbortError } from '@aztec/foundation/error';
-import { Fr } from '@aztec/foundation/fields';
 import { createLogger } from '@aztec/foundation/log';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
 import { assertLength } from '@aztec/foundation/serialize';
 import { pushTestData } from '@aztec/foundation/testing';
 import { elapsed } from '@aztec/foundation/timer';
 import type { TreeNodeLocation } from '@aztec/foundation/trees';
-import { readAvmMinimalPublicTxInputsFromFile } from '@aztec/simulator/public/fixtures';
-import { EthAddress, createBlockEndMarker } from '@aztec/stdlib/block';
+import { EthAddress } from '@aztec/stdlib/block';
 import type {
   EpochProver,
   ForkMerkleTreeOperations,
   MerkleTreeWriteOperations,
   PublicInputsAndRecursiveProof,
+  ReadonlyWorldStateAccess,
   ServerCircuitProver,
 } from '@aztec/stdlib/interfaces/server';
 import type { Proof } from '@aztec/stdlib/proofs';
@@ -54,7 +55,6 @@ import {
 import { inspect } from 'util';
 
 import {
-  buildBlockHeaderFromTxs,
   buildHeaderFromCircuitOutputs,
   getLastSiblingPath,
   getPublicChonkVerifierPrivateInputsFromTx,
@@ -73,6 +73,11 @@ import { TxProvingState } from './tx-proving-state.js';
 
 const logger = createLogger('prover-client:orchestrator');
 
+type WorldStateFork = {
+  fork: MerkleTreeWriteOperations;
+  cleanupPromise: Promise<void> | undefined;
+};
+
 /**
  * Implements an event driven proving scheduler to build the recursive proof tree. The idea being:
  * 1. Transactions are provided to the scheduler post simulation.
@@ -93,12 +98,14 @@ export class ProvingOrchestrator implements EpochProver {
 
   private provingPromise: Promise<ProvingResult> | undefined = undefined;
   private metrics: ProvingOrchestratorMetrics;
-  private dbs: Map<number, MerkleTreeWriteOperations> = new Map();
+  // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+  private dbs: Map<BlockNumber, WorldStateFork> = new Map();
 
   constructor(
-    private dbProvider: ForkMerkleTreeOperations,
+    private dbProvider: ReadonlyWorldStateAccess & ForkMerkleTreeOperations,
     private prover: ServerCircuitProver,
     private readonly proverId: EthAddress,
+    private readonly cancelJobsOnStop: boolean = false,
     telemetryClient: TelemetryClient = getTelemetryClient(),
   ) {
     this.metrics = new ProvingOrchestratorMetrics(telemetryClient, 'ProvingOrchestrator');
@@ -112,13 +119,17 @@ export class ProvingOrchestrator implements EpochProver {
     return this.proverId;
   }
 
+  public getNumActiveForks() {
+    return this.dbs.size;
+  }
+
   public stop(): Promise<void> {
     this.cancel();
     return Promise.resolve();
   }
 
   public startNewEpoch(
-    epochNumber: number,
+    epochNumber: EpochNumber,
     totalNumCheckpoints: number,
     finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
   ) {
@@ -142,12 +153,19 @@ export class ProvingOrchestrator implements EpochProver {
     this.provingPromise = promise;
   }
 
+  /**
+   * Starts a new checkpoint.
+   * @param checkpointIndex - The index of the checkpoint in the epoch.
+   * @param constants - The constants for this checkpoint.
+   * @param l1ToL2Messages - The set of L1 to L2 messages to be inserted at the beginning of this checkpoint.
+   * @param totalNumBlocks - The total number of blocks expected in the checkpoint (must be at least one).
+   * @param headerOfLastBlockInPreviousCheckpoint - The header of the last block in the previous checkpoint.
+   */
   public async startNewCheckpoint(
     checkpointIndex: number,
     constants: CheckpointConstantData,
     l1ToL2Messages: Fr[],
     totalNumBlocks: number,
-    totalNumBlobFields: number,
     headerOfLastBlockInPreviousCheckpoint: BlockHeader,
   ) {
     if (!this.provingState) {
@@ -162,8 +180,8 @@ export class ProvingOrchestrator implements EpochProver {
     const lastBlockNumber = headerOfLastBlockInPreviousCheckpoint.globalVariables.blockNumber;
     const db = await this.dbProvider.fork(lastBlockNumber);
 
-    const firstBlockNumber = lastBlockNumber + 1;
-    this.dbs.set(firstBlockNumber, db);
+    const firstBlockNumber = BlockNumber(lastBlockNumber + 1);
+    this.dbs.set(firstBlockNumber, { fork: db, cleanupPromise: undefined });
 
     // Get archive sibling path before any block in this checkpoint lands.
     const lastArchiveSiblingPath = await getLastSiblingPath(MerkleTreeId.ARCHIVE, db);
@@ -180,7 +198,6 @@ export class ProvingOrchestrator implements EpochProver {
       checkpointIndex,
       constants,
       totalNumBlocks,
-      totalNumBlobFields,
       headerOfLastBlockInPreviousCheckpoint,
       lastArchiveSiblingPath,
       l1ToL2Messages,
@@ -201,7 +218,7 @@ export class ProvingOrchestrator implements EpochProver {
   @trackSpan('ProvingOrchestrator.startNewBlock', blockNumber => ({
     [Attributes.BLOCK_NUMBER]: blockNumber,
   }))
-  public async startNewBlock(blockNumber: number, timestamp: UInt64, totalNumTxs: number) {
+  public async startNewBlock(blockNumber: BlockNumber, timestamp: UInt64, totalNumTxs: number) {
     if (!this.provingState) {
       throw new Error('Empty epoch proving state. Call startNewEpoch before starting a block.');
     }
@@ -216,21 +233,21 @@ export class ProvingOrchestrator implements EpochProver {
     }
 
     const constants = checkpointProvingState.constants;
-    logger.info(`Starting block ${blockNumber} for slot ${constants.slotNumber.toNumber()}.`);
+    logger.info(`Starting block ${blockNumber} for slot ${constants.slotNumber}.`);
 
     // Fork the db only when it's not already set. The db for the first block is set in `startNewCheckpoint`.
     if (!this.dbs.has(blockNumber)) {
       // Fork world state at the end of the immediately previous block
-      const db = await this.dbProvider.fork(blockNumber - 1);
-      this.dbs.set(blockNumber, db);
+      const db = await this.dbProvider.fork(BlockNumber(blockNumber - 1));
+      this.dbs.set(blockNumber, { fork: db, cleanupPromise: undefined });
     }
-    const db = this.dbs.get(blockNumber)!;
+    const db = this.dbs.get(blockNumber)!.fork;
 
     // Get archive snapshot and sibling path before any txs in this block lands.
     const lastArchiveTreeSnapshot = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
     const lastArchiveSiblingPath = await getRootTreeSiblingPath(MerkleTreeId.ARCHIVE, db);
 
-    const blockProvingState = await checkpointProvingState.startNewBlock(
+    const blockProvingState = checkpointProvingState.startNewBlock(
       blockNumber,
       timestamp,
       totalNumTxs,
@@ -248,11 +265,16 @@ export class ProvingOrchestrator implements EpochProver {
     // Because `addTxs` won't be called for a block without txs, and that's where the sponge blob state is computed.
     // We need to set its end sponge blob here, which will become the start sponge blob for the next block.
     if (totalNumTxs === 0) {
+      const endState = await db.getStateReference();
+      blockProvingState.setEndState(endState);
+
       const endSpongeBlob = blockProvingState.getStartSpongeBlob().clone();
-      await endSpongeBlob.absorb([createBlockEndMarker(0)]);
+      const blockEndBlobFields = blockProvingState.getBlockEndBlobFields();
+      await endSpongeBlob.absorb(blockEndBlobFields);
       blockProvingState.setEndSpongeBlob(endSpongeBlob);
 
-      // And also try to accumulate the blobs as far as we can:
+      // Try to accumulate the out hashes and blobs as far as we can:
+      await this.provingState.accumulateCheckpointOutHashes();
       await this.provingState.setBlobAccumulators();
     }
   }
@@ -276,7 +298,7 @@ export class ProvingOrchestrator implements EpochProver {
       return;
     }
 
-    const blockNumber = txs[0].globalVariables.blockNumber;
+    const blockNumber = BlockNumber(txs[0].globalVariables.blockNumber);
     const provingState = this.provingState.getBlockProvingStateByBlockNumber(blockNumber!);
     if (!provingState) {
       throw new Error(`Proving state for block ${blockNumber} not found. Call startNewBlock first.`);
@@ -294,7 +316,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     logger.info(`Adding ${txs.length} transactions to block ${blockNumber}`);
 
-    const db = this.dbs.get(blockNumber)!;
+    const db = this.dbs.get(blockNumber)!.fork;
     const lastArchive = provingState.lastArchiveTreeSnapshot;
     const newL1ToL2MessageTreeSnapshot = provingState.newL1ToL2MessageTreeSnapshot;
     const spongeBlobState = provingState.getStartSpongeBlob().clone();
@@ -307,7 +329,7 @@ export class ProvingOrchestrator implements EpochProver {
 
         validateTx(tx);
 
-        logger.info(`Received transaction: ${tx.hash}`);
+        logger.debug(`Received transaction: ${tx.hash}`);
 
         const startSpongeBlob = spongeBlobState.clone();
         const [hints, treeSnapshots] = await this.prepareBaseRollupInputs(
@@ -341,11 +363,16 @@ export class ProvingOrchestrator implements EpochProver {
       }
     }
 
-    await spongeBlobState.absorb([createBlockEndMarker(txs.length)]);
+    const endState = await db.getStateReference();
+    provingState.setEndState(endState);
+
+    const blockEndBlobFields = provingState.getBlockEndBlobFields();
+    await spongeBlobState.absorb(blockEndBlobFields);
 
     provingState.setEndSpongeBlob(spongeBlobState);
 
-    // Txs have been added to the block. Now try to accumulate the blobs as far as we can:
+    // Txs have been added to the block. Now try to accumulate the out hashes and blobs as far as we can:
+    await this.provingState.accumulateCheckpointOutHashes();
     await this.provingState.setBlobAccumulators();
   }
 
@@ -382,10 +409,10 @@ export class ProvingOrchestrator implements EpochProver {
    * Marks the block as completed.
    * Computes the block header and updates the archive tree.
    */
-  @trackSpan('ProvingOrchestrator.setBlockCompleted', (blockNumber: number) => ({
+  @trackSpan('ProvingOrchestrator.setBlockCompleted', (blockNumber: BlockNumber) => ({
     [Attributes.BLOCK_NUMBER]: blockNumber,
   }))
-  public async setBlockCompleted(blockNumber: number, expectedHeader?: BlockHeader): Promise<BlockHeader> {
+  public async setBlockCompleted(blockNumber: BlockNumber, expectedHeader?: BlockHeader): Promise<BlockHeader> {
     const provingState = this.provingState?.getBlockProvingStateByBlockNumber(blockNumber);
     if (!provingState) {
       throw new Error(`Block proving state for ${blockNumber} not found`);
@@ -408,39 +435,25 @@ export class ProvingOrchestrator implements EpochProver {
       );
     }
 
-    // And build the block header
+    // Given we've applied every change from this block, now assemble the block header:
     logger.verbose(`Block ${blockNumber} completed. Assembling header.`);
-    const header = await this.buildL2BlockHeader(provingState, expectedHeader);
-
-    await this.verifyBuiltBlockAgainstSyncedState(provingState);
-
-    return header;
-  }
-
-  private async buildL2BlockHeader(provingState: BlockProvingState, expectedHeader?: BlockHeader) {
-    // Collect all txs in this block to build the header. The function calling this has made sure that all txs have been added.
-    const txs = provingState.getProcessedTxs();
-
-    const startSpongeBlob = provingState.getStartSpongeBlob();
-
-    // Get db for this block
-    const db = this.dbs.get(provingState.blockNumber)!;
-
-    // Given we've applied every change from this block, now assemble the block header
-    // and update the archive tree, so we're ready to start processing the next block
-    const header = await buildBlockHeaderFromTxs(txs, provingState.getGlobalVariables(), startSpongeBlob, db);
+    const header = await provingState.buildBlockHeader();
 
     if (expectedHeader && !header.equals(expectedHeader)) {
       logger.error(`Block header mismatch: header=${header} expectedHeader=${expectedHeader}`);
       throw new Error('Block header mismatch');
     }
 
+    // Get db for this block
+    const db = this.dbs.get(provingState.blockNumber)!.fork;
+
+    // Update the archive tree, so we're ready to start processing the next block:
     logger.verbose(
       `Updating archive tree with block ${provingState.blockNumber} header ${(await header.hash()).toString()}`,
     );
     await db.updateArchive(header);
 
-    provingState.setBuiltBlockHeader(header);
+    await this.verifyBuiltBlockAgainstSyncedState(provingState);
 
     return header;
   }
@@ -468,7 +481,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     // Get db for this block
     const blockNumber = provingState.blockNumber;
-    const db = this.dbs.get(blockNumber)!;
+    const db = this.dbs.get(blockNumber)!.fork;
 
     const newArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
     const syncedArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, this.dbProvider.getSnapshot(blockNumber));
@@ -493,20 +506,19 @@ export class ProvingOrchestrator implements EpochProver {
     // is aborted and never reaches this point, it will leak the fork. We need to add a global cleanup,
     // but have to make sure it only runs once all operations are completed, otherwise some function here
     // will attempt to access the fork after it was closed.
-    logger.debug(`Cleaning up world state fork for ${blockNumber}`);
-    void this.dbs
-      .get(blockNumber)
-      ?.close()
-      .then(() => this.dbs.delete(blockNumber))
-      .catch(err => logger.error(`Error closing db for block ${blockNumber}`, err));
+    void this.cleanupDBFork(blockNumber);
   }
 
   /**
-   * Cancel any further proving
+   * Cancel any further proving.
+   * If cancelJobsOnStop is true, aborts all pending jobs with the broker (which marks them as 'Aborted').
+   * If cancelJobsOnStop is false (default), jobs remain in the broker queue and can be reused on restart/reorg.
    */
   public cancel() {
-    for (const controller of this.pendingProvingJobs) {
-      controller.abort();
+    if (this.cancelJobsOnStop) {
+      for (const controller of this.pendingProvingJobs) {
+        controller.abort();
+      }
     }
 
     this.provingState?.cancel();
@@ -541,6 +553,24 @@ export class ProvingOrchestrator implements EpochProver {
     return epochProofResult;
   }
 
+  private async cleanupDBFork(blockNumber: BlockNumber): Promise<void> {
+    logger.debug(`Cleaning up world state fork for ${blockNumber}`);
+    const fork = this.dbs.get(blockNumber);
+    if (!fork) {
+      return;
+    }
+
+    try {
+      if (!fork.cleanupPromise) {
+        fork.cleanupPromise = fork.fork.close();
+      }
+      await fork.cleanupPromise;
+      this.dbs.delete(blockNumber);
+    } catch (err) {
+      logger.error(`Error closing db for block ${blockNumber}`, err);
+    }
+  }
+
   /**
    * Enqueue a job to be scheduled
    * @param provingState - The proving state object being operated on
@@ -858,19 +888,22 @@ export class ProvingOrchestrator implements EpochProver {
         },
       ),
       async result => {
-        // If the proofs were slower than the block header building, then we need to try validating the block header hashes here.
-        await this.verifyBuiltBlockAgainstSyncedState(provingState);
-
         logger.debug(`Completed ${rollupType} proof for block ${provingState.blockNumber}`);
 
         const leafLocation = provingState.setBlockRootRollupProof(result);
         const checkpointProvingState = provingState.parentCheckpoint;
 
+        // If the proofs were slower than the block header building, then we need to try validating the block header hashes here.
+        await this.verifyBuiltBlockAgainstSyncedState(provingState);
+
         if (checkpointProvingState.totalNumBlocks === 1) {
           this.checkAndEnqueueCheckpointRootRollup(checkpointProvingState);
         } else {
           this.checkAndEnqueueNextBlockMergeRollup(checkpointProvingState, leafLocation);
         }
+
+        // We are finished with the block at this point, ensure the fork is cleaned up
+        void this.cleanupDBFork(provingState.blockNumber);
       },
     );
   }
@@ -1214,8 +1247,6 @@ export class ProvingOrchestrator implements EpochProver {
 
     const txProvingState = provingState.getTxProvingState(txIndex);
 
-    // This function tries to do AVM proving. If there is a failure, it fakes the proof unless AVM_PROVING_STRICT is defined.
-    // Nothing downstream depends on the AVM proof yet. So having this mode lets us incrementally build the AVM circuit.
     const doAvmProving = wrapCallbackInSpan(
       this.tracer,
       'ProvingOrchestrator.prover.getAvmProof',
@@ -1224,36 +1255,13 @@ export class ProvingOrchestrator implements EpochProver {
       },
       async (signal: AbortSignal) => {
         const inputs = txProvingState.getAvmInputs();
-        try {
-          // TODO(#14234)[Unconditional PIs validation]: Remove the whole try-catch logic and
-          // just keep the next line but removing the second argument (false).
-          return await this.prover.getAvmProof(inputs, false, signal, provingState.epochNumber);
-        } catch (err) {
-          if (process.env.AVM_PROVING_STRICT) {
-            logger.error(`Error thrown when proving AVM circuit with AVM_PROVING_STRICT on`, err);
-            throw err;
-          } else {
-            logger.warn(
-              `Error thrown when proving AVM circuit but AVM_PROVING_STRICT is off. Use snapshotted
-               AVM inputs and carrying on. ${inspect(err)}.`,
-            );
-
-            try {
-              this.metrics.incAvmFallback();
-              const snapshotAvmPrivateInputs = readAvmMinimalPublicTxInputsFromFile();
-              return await this.prover.getAvmProof(snapshotAvmPrivateInputs, true, signal, provingState.epochNumber);
-            } catch (err) {
-              logger.error(`Error thrown when proving snapshotted AVM inputs.`, err);
-              throw err;
-            }
-          }
-        }
+        return await this.prover.getAvmProof(inputs, signal, provingState.epochNumber);
       },
     );
 
-    this.deferredProving(provingState, doAvmProving, proofAndVk => {
+    this.deferredProving(provingState, doAvmProving, proof => {
       logger.debug(`Proven VM for tx index: ${txIndex}`);
-      txProvingState.setAvmProof(proofAndVk);
+      txProvingState.setAvmProof(proof);
       this.checkAndEnqueueBaseRollup(provingState, txIndex);
     });
   }

package/src/orchestrator/orchestrator_metrics.ts

@@ -1,41 +1,18 @@
-import {
-  type Histogram,
-  Metrics,
-  type TelemetryClient,
-  type Tracer,
-  type UpDownCounter,
-  ValueType,
-} from '@aztec/telemetry-client';
+import { type Histogram, Metrics, type TelemetryClient, type Tracer } from '@aztec/telemetry-client';
 
 export class ProvingOrchestratorMetrics {
   public readonly tracer: Tracer;
 
   private baseRollupInputsDuration: Histogram;
-  private avmFallbackCount: UpDownCounter;
 
   constructor(client: TelemetryClient, name = 'ProvingOrchestrator') {
     this.tracer = client.getTracer(name);
     const meter = client.getMeter(name);
 
-    this.baseRollupInputsDuration = meter.createHistogram(Metrics.PROVING_ORCHESTRATOR_BASE_ROLLUP_INPUTS_DURATION, {
-      unit: 'ms',
-      description: 'Duration to build base rollup inputs',
-      valueType: ValueType.INT,
-    });
-
-    this.avmFallbackCount = meter.createUpDownCounter(Metrics.PROVING_ORCHESTRATOR_AVM_FALLBACK_COUNT, {
-      description: 'How many times the AVM fallback was used',
-      valueType: ValueType.INT,
-    });
-
-    this.avmFallbackCount.add(0);
+    this.baseRollupInputsDuration = meter.createHistogram(Metrics.PROVING_ORCHESTRATOR_BASE_ROLLUP_INPUTS_DURATION);
   }
 
   recordBaseRollupInputs(durationMs: number) {
     this.baseRollupInputsDuration.record(Math.ceil(durationMs));
   }
-
-  incAvmFallback() {
-    this.avmFallbackCount.add(1);
-  }
 }

package/src/orchestrator/tx-proving-state.ts

@@ -1,14 +1,9 @@
-import {
-  AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED,
-  AVM_VK_INDEX,
-  NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
-} from '@aztec/constants';
-import type { Fr } from '@aztec/foundation/fields';
+import { AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED, NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH } from '@aztec/constants';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import { getVkData } from '@aztec/noir-protocol-circuits-types/server/vks';
-import { getVKSiblingPath } from '@aztec/noir-protocol-circuits-types/vk-tree';
 import type { AvmCircuitInputs } from '@aztec/stdlib/avm';
-import type { ProofAndVerificationKey, PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
-import { ProofData } from '@aztec/stdlib/proofs';
+import type { PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
+import { ProofData, ProofDataForFixedVk, RecursiveProof } from '@aztec/stdlib/proofs';
 import {
   type BaseRollupHints,
   PrivateBaseRollupHints,
@@ -20,7 +15,6 @@ import {
 import type { CircuitName } from '@aztec/stdlib/stats';
 import type { AppendOnlyTreeSnapshot, MerkleTreeId } from '@aztec/stdlib/trees';
 import type { ProcessedTx } from '@aztec/stdlib/tx';
-import { VerificationKeyData, VkData } from '@aztec/stdlib/vks';
 
 import {
   getChonkProofFromTx,
@@ -38,7 +32,7 @@ export class TxProvingState {
     PublicChonkVerifierPublicInputs,
     typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
   >;
-  private avm?: ProofAndVerificationKey<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>;
+  private avmProof?: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>;
 
   constructor(
     public readonly processedTx: ProcessedTx,
@@ -52,7 +46,7 @@ export class TxProvingState {
   }
 
   public ready() {
-    return !this.requireAvmProof || (!!this.avm && !!this.publicChonkVerifier);
+    return !this.requireAvmProof || (!!this.avmProof && !!this.publicChonkVerifier);
   }
 
   public getAvmInputs(): AvmCircuitInputs {
@@ -86,8 +80,8 @@ export class TxProvingState {
     this.publicChonkVerifier = publicChonkVerifierProofAndVk;
   }
 
-  public setAvmProof(avmProofAndVk: ProofAndVerificationKey<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>) {
-    this.avm = avmProofAndVk;
+  public setAvmProof(avmProof: RecursiveProof<typeof AVM_V2_PROOF_LENGTH_IN_FIELDS_PADDED>) {
+    this.avmProof = avmProof;
   }
 
   #getPrivateBaseInputs() {
@@ -111,7 +105,7 @@ export class TxProvingState {
     if (!this.publicChonkVerifier) {
       throw new Error('Tx not ready for proving base rollup: public chonk verifier proof undefined');
     }
-    if (!this.avm) {
+    if (!this.avmProof) {
       throw new Error('Tx not ready for proving base rollup: avm proof undefined');
     }
     if (!(this.baseRollupHints instanceof PublicBaseRollupHints)) {
@@ -120,19 +114,8 @@ export class TxProvingState {
 
     const publicChonkVerifierProofData = toProofData(this.publicChonkVerifier);
 
-    const avmProofData = new ProofData(
-      this.processedTx.avmProvingRequest.inputs.publicInputs,
-      this.avm.proof,
-      this.#getVkData(this.avm!.verificationKey, AVM_VK_INDEX),
-    );
+    const avmProofData = new ProofDataForFixedVk(this.processedTx.avmProvingRequest.inputs.publicInputs, this.avmProof);
 
     return new PublicTxBaseRollupPrivateInputs(publicChonkVerifierProofData, avmProofData, this.baseRollupHints);
   }
-
-  #getVkData(verificationKey: VerificationKeyData, vkIndex: number) {
-    // TODO(#17162): Add avm vk hash to the tree and call `getVkData('AVM')` instead.
-    // Below will return a path to an empty leaf.
-    const vkPath = getVKSiblingPath(vkIndex);
-    return new VkData(verificationKey, vkIndex, vkPath);
-  }
 }

package/src/prover-client/factory.ts

@@ -1,4 +1,8 @@
-import type { ForkMerkleTreeOperations, ProvingJobBroker } from '@aztec/stdlib/interfaces/server';
+import type {
+  ForkMerkleTreeOperations,
+  ProvingJobBroker,
+  ReadonlyWorldStateAccess,
+} from '@aztec/stdlib/interfaces/server';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import type { ProverClientConfig } from '../config.js';
@@ -6,7 +10,7 @@ import { ProverClient } from './prover-client.js';
 
 export function createProverClient(
   config: ProverClientConfig,
-  worldState: ForkMerkleTreeOperations,
+  worldState: ForkMerkleTreeOperations & ReadonlyWorldStateAccess,
   broker: ProvingJobBroker,
   telemetry: TelemetryClient = getTelemetryClient(),
 ) {

package/src/prover-client/prover-client.ts

@@ -11,6 +11,7 @@ import {
   type ProvingJobBroker,
   type ProvingJobConsumer,
   type ProvingJobProducer,
+  type ReadonlyWorldStateAccess,
   type ServerCircuitProver,
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
@@ -33,7 +34,7 @@ export class ProverClient implements EpochProverManager {
 
   private constructor(
     private config: ProverClientConfig,
-    private worldState: ForkMerkleTreeOperations,
+    private worldState: ForkMerkleTreeOperations & ReadonlyWorldStateAccess,
     private orchestratorClient: ProvingJobProducer,
     private agentClient?: ProvingJobConsumer,
     private telemetry: TelemetryClient = getTelemetryClient(),
@@ -45,7 +46,13 @@ export class ProverClient implements EpochProverManager {
 
   public createEpochProver(): EpochProver {
     const facade = new BrokerCircuitProverFacade(this.orchestratorClient, this.proofStore, this.failedProofStore);
-    const orchestrator = new ProvingOrchestrator(this.worldState, facade, this.config.proverId, this.telemetry);
+    const orchestrator = new ProvingOrchestrator(
+      this.worldState,
+      facade,
+      this.config.proverId,
+      this.config.cancelJobsOnStop,
+      this.telemetry,
+    );
     return new ServerEpochProver(facade, orchestrator);
   }
 
@@ -99,7 +106,7 @@ export class ProverClient implements EpochProverManager {
    */
   public static async new(
     config: ProverClientConfig,
-    worldState: ForkMerkleTreeOperations,
+    worldState: ForkMerkleTreeOperations & ReadonlyWorldStateAccess,
     broker: ProvingJobBroker,
     telemetry: TelemetryClient = getTelemetryClient(),
   ) {
@@ -129,15 +136,7 @@ export class ProverClient implements EpochProverManager {
     const prover = await buildServerCircuitProver(this.config, this.telemetry);
     this.agents = times(
       this.config.proverAgentCount,
-      () =>
-        new ProvingAgent(
-          this.agentClient!,
-          proofStore,
-          prover,
-          [],
-          this.config.proverAgentPollIntervalMs,
-          this.telemetry,
-        ),
+      () => new ProvingAgent(this.agentClient!, proofStore, prover, [], this.config.proverAgentPollIntervalMs),
     );
 
     await Promise.all(this.agents.map(agent => agent.start()));

package/src/prover-client/server-epoch-prover.ts

@@ -1,5 +1,6 @@
-import type { BatchedBlob, FinalBlobBatchingChallenges } from '@aztec/blob-lib';
-import type { Fr } from '@aztec/foundation/fields';
+import type { BatchedBlob, FinalBlobBatchingChallenges } from '@aztec/blob-lib/types';
+import { BlockNumber, EpochNumber } from '@aztec/foundation/branded-types';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EthAddress } from '@aztec/stdlib/block';
 import type { EpochProver } from '@aztec/stdlib/interfaces/server';
 import type { Proof } from '@aztec/stdlib/proofs';
@@ -18,7 +19,7 @@ export class ServerEpochProver implements EpochProver {
   ) {}
 
   startNewEpoch(
-    epochNumber: number,
+    epochNumber: EpochNumber,
     totalNumCheckpoints: number,
     finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
   ): void {
@@ -30,7 +31,6 @@ export class ServerEpochProver implements EpochProver {
     constants: CheckpointConstantData,
     l1ToL2Messages: Fr[],
     totalNumBlocks: number,
-    totalNumBlobFields: number,
     headerOfLastBlockInPreviousCheckpoint: BlockHeader,
   ): Promise<void> {
     return this.orchestrator.startNewCheckpoint(
@@ -38,14 +38,13 @@ export class ServerEpochProver implements EpochProver {
       constants,
       l1ToL2Messages,
       totalNumBlocks,
-      totalNumBlobFields,
       headerOfLastBlockInPreviousCheckpoint,
     );
   }
   startChonkVerifierCircuits(txs: Tx[]): Promise<void> {
     return this.orchestrator.startChonkVerifierCircuits(txs);
   }
-  setBlockCompleted(blockNumber: number, expectedBlockHeader?: BlockHeader): Promise<BlockHeader> {
+  setBlockCompleted(blockNumber: BlockNumber, expectedBlockHeader?: BlockHeader): Promise<BlockHeader> {
     return this.orchestrator.setBlockCompleted(blockNumber, expectedBlockHeader);
   }
   finalizeEpoch(): Promise<{ publicInputs: RootRollupPublicInputs; proof: Proof; batchedBlobInputs: BatchedBlob }> {
@@ -61,7 +60,7 @@ export class ServerEpochProver implements EpochProver {
     await this.facade.stop();
     await this.orchestrator.stop();
   }
-  startNewBlock(blockNumber: number, timestamp: UInt64, totalNumTxs: number): Promise<void> {
+  startNewBlock(blockNumber: BlockNumber, timestamp: UInt64, totalNumTxs: number): Promise<void> {
     return this.orchestrator.startNewBlock(blockNumber, timestamp, totalNumTxs);
   }
   addTxs(txs: ProcessedTx[]): Promise<void> {