@aztec/prover-client 0.63.1 → 0.65.0

package/src/mocks/test_context.ts

@@ -13,10 +13,8 @@ import { type Fr } from '@aztec/foundation/fields';
 import { type DebugLogger } from '@aztec/foundation/log';
 import { openTmpStore } from '@aztec/kv-store/utils';
 import {
-  type PublicExecutionResult,
-  PublicExecutionResultBuilder,
-  type PublicExecutor,
   PublicProcessor,
+  PublicTxSimulator,
   type SimulationProvider,
   WASMSimulator,
   type WorldStateDB,
@@ -25,10 +23,12 @@ import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 import { MerkleTrees } from '@aztec/world-state';
 import { NativeWorldStateService } from '@aztec/world-state/native';
 
+import { jest } from '@jest/globals';
 import * as fs from 'fs/promises';
 import { type MockProxy, mock } from 'jest-mock-extended';
 
 import { TestCircuitProver } from '../../../bb-prover/src/test/test_circuit_prover.js';
+import { AvmFinalizedCallResult } from '../../../simulator/src/avm/avm_contract_call_result.js';
 import { type AvmPersistableStateManager } from '../../../simulator/src/avm/journal/journal.js';
 import { ProvingOrchestrator } from '../orchestrator/index.js';
 import { MemoryProvingQueue } from '../prover-agent/memory-proving-queue.js';
@@ -37,7 +37,7 @@ import { getEnvironmentConfig, getSimulationProvider, makeGlobals } from './fixt
 
 export class TestContext {
   constructor(
-    public publicExecutor: MockProxy<PublicExecutor>,
+    public publicTxSimulator: PublicTxSimulator,
     public worldStateDB: MockProxy<WorldStateDB>,
     public publicProcessor: PublicProcessor,
     public simulationProvider: SimulationProvider,
@@ -66,7 +66,6 @@ export class TestContext {
     const directoriesToCleanup: string[] = [];
     const globalVariables = makeGlobals(blockNumber);
 
-    const publicExecutor = mock<PublicExecutor>();
     const worldStateDB = mock<WorldStateDB>();
     const telemetry = new NoopTelemetryClient();
 
@@ -83,13 +82,15 @@ export class TestContext {
       publicDb = await ws.getLatest();
       proverDb = await ws.getLatest();
     }
+    worldStateDB.getMerkleInterface.mockReturnValue(publicDb);
 
-    const processor = PublicProcessor.create(
+    const publicTxSimulator = new PublicTxSimulator(publicDb, worldStateDB, telemetry, globalVariables);
+    const processor = new PublicProcessor(
       publicDb,
-      publicExecutor,
       globalVariables,
       Header.empty(),
       worldStateDB,
+      publicTxSimulator,
       telemetry,
     );
 
@@ -124,7 +125,7 @@ export class TestContext {
     agent.start(queue);
 
     return new this(
-      publicExecutor,
+      publicTxSimulator,
       worldStateDB,
       processor,
       simulationProvider,
@@ -154,25 +155,24 @@ export class TestContext {
   ) {
     const defaultExecutorImplementation = (
       _stateManager: AvmPersistableStateManager,
-      execution: PublicExecutionRequest,
-      _globalVariables: GlobalVariables,
+      executionRequest: PublicExecutionRequest,
       allocatedGas: Gas,
-      _transactionFee?: Fr,
+      _transactionFee: Fr,
+      _fnName: string,
     ) => {
       for (const tx of txs) {
         const allCalls = tx.publicTeardownFunctionCall.isEmpty()
           ? tx.enqueuedPublicFunctionCalls
           : [...tx.enqueuedPublicFunctionCalls, tx.publicTeardownFunctionCall];
         for (const request of allCalls) {
-          if (execution.callContext.equals(request.callContext)) {
-            const result = PublicExecutionResultBuilder.empty().build({
-              endGasLeft: allocatedGas,
-            });
-            return Promise.resolve(result);
+          if (executionRequest.callContext.equals(request.callContext)) {
+            return Promise.resolve(
+              new AvmFinalizedCallResult(/*reverted=*/ false, /*output=*/ [], /*gasLeft=*/ allocatedGas),
+            );
           }
         }
       }
-      throw new Error(`Unexpected execution request: ${execution}`);
+      throw new Error(`Unexpected execution request: ${executionRequest}`);
     };
     return await this.processPublicFunctionsWithMockExecutorImplementation(
       txs,
@@ -183,21 +183,36 @@ export class TestContext {
     );
   }
 
-  public async processPublicFunctionsWithMockExecutorImplementation(
+  private async processPublicFunctionsWithMockExecutorImplementation(
     txs: Tx[],
     maxTransactions: number,
     txHandler?: ProcessedTxHandler,
     txValidator?: TxValidator<ProcessedTx>,
     executorMock?: (
       stateManager: AvmPersistableStateManager,
-      execution: PublicExecutionRequest,
-      globalVariables: GlobalVariables,
+      executionRequest: PublicExecutionRequest,
       allocatedGas: Gas,
-      transactionFee?: Fr,
-    ) => Promise<PublicExecutionResult>,
+      transactionFee: Fr,
+      fnName: string,
+    ) => Promise<AvmFinalizedCallResult>,
   ) {
+    // Mock the internal private function. Borrowed from https://stackoverflow.com/a/71033167
+    const simulateInternal: jest.SpiedFunction<
+      (
+        stateManager: AvmPersistableStateManager,
+        executionResult: any,
+        allocatedGas: Gas,
+        transactionFee: any,
+        fnName: any,
+      ) => Promise<AvmFinalizedCallResult>
+    > = jest.spyOn(
+      this.publicTxSimulator as unknown as {
+        simulateEnqueuedCallInternal: PublicTxSimulator['simulateEnqueuedCallInternal'];
+      },
+      'simulateEnqueuedCallInternal',
+    );
     if (executorMock) {
-      this.publicExecutor.simulate.mockImplementation(executorMock);
+      simulateInternal.mockImplementation(executorMock);
     }
     return await this.publicProcessor.process(txs, maxTransactions, txHandler, txValidator);
   }

package/src/orchestrator/block-building-helpers.ts

@@ -10,7 +10,6 @@ import {
   ARCHIVE_HEIGHT,
   AppendOnlyTreeSnapshot,
   type BaseOrMergeRollupPublicInputs,
-  BaseRollupHints,
   BlockMergeRollupInputs,
   type BlockRootOrBlockMergePublicInputs,
   ConstantRollupData,
@@ -32,25 +31,25 @@ import {
   NULLIFIER_TREE_HEIGHT,
   NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
   NullifierLeafPreimage,
-  PUBLIC_DATA_SUBTREE_HEIGHT,
-  PUBLIC_DATA_SUBTREE_SIBLING_PATH_LENGTH,
   PUBLIC_DATA_TREE_HEIGHT,
   type ParityPublicInputs,
   PartialStateReference,
   PreviousRollupBlockData,
   PreviousRollupData,
+  PrivateBaseRollupHints,
+  PrivateBaseStateDiffHints,
+  PublicBaseRollupHints,
+  PublicBaseStateDiffHints,
   PublicDataHint,
   PublicDataTreeLeaf,
-  type PublicDataTreeLeafPreimage,
-  PublicDataWrite,
+  PublicDataTreeLeafPreimage,
   type RecursiveProof,
   RootRollupInputs,
-  StateDiffHints,
   StateReference,
   VK_TREE_HEIGHT,
   type VerificationKeyAsFields,
 } from '@aztec/circuits.js';
-import { assertPermutation, makeTuple } from '@aztec/foundation/array';
+import { makeTuple } from '@aztec/foundation/array';
 import { padArrayEnd } from '@aztec/foundation/collection';
 import { sha256Trunc } from '@aztec/foundation/crypto';
 import { type DebugLogger } from '@aztec/foundation/log';
@@ -123,6 +122,7 @@ export async function buildBaseRollupHints(
     padArrayEnd(tx.txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX).map(n => n.toBuffer()),
     NULLIFIER_SUBTREE_HEIGHT,
   );
+
   if (nullifierWitnessLeaves === undefined) {
     throw new Error(`Could not craft nullifier batch insertion proofs`);
   }
@@ -139,45 +139,109 @@ export async function buildBaseRollupHints(
     i < nullifierSubtreeSiblingPathArray.length ? nullifierSubtreeSiblingPathArray[i] : Fr.ZERO,
   );
 
-  const publicDataSiblingPath = txPublicDataUpdateRequestInfo.newPublicDataSubtreeSiblingPath;
-
-  const stateDiffHints = StateDiffHints.from({
-    nullifierPredecessorPreimages: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
-      i < nullifierWitnessLeaves.length
-        ? (nullifierWitnessLeaves[i].leafPreimage as NullifierLeafPreimage)
-        : NullifierLeafPreimage.empty(),
-    ),
-    nullifierPredecessorMembershipWitnesses: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
-      i < nullifierPredecessorMembershipWitnessesWithoutPadding.length
-        ? nullifierPredecessorMembershipWitnessesWithoutPadding[i]
-        : makeEmptyMembershipWitness(NULLIFIER_TREE_HEIGHT),
-    ),
-    sortedNullifiers: makeTuple(MAX_NULLIFIERS_PER_TX, i => Fr.fromBuffer(sortednullifiers[i])),
-    sortedNullifierIndexes: makeTuple(MAX_NULLIFIERS_PER_TX, i => sortedNewLeavesIndexes[i]),
-    noteHashSubtreeSiblingPath,
-    nullifierSubtreeSiblingPath,
-    publicDataSiblingPath,
-  });
+  if (tx.avmProvingRequest) {
+    // Build public base rollup hints
+    const stateDiffHints = PublicBaseStateDiffHints.from({
+      nullifierPredecessorPreimages: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
+        i < nullifierWitnessLeaves.length
+          ? (nullifierWitnessLeaves[i].leafPreimage as NullifierLeafPreimage)
+          : NullifierLeafPreimage.empty(),
+      ),
+      nullifierPredecessorMembershipWitnesses: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
+        i < nullifierPredecessorMembershipWitnessesWithoutPadding.length
+          ? nullifierPredecessorMembershipWitnessesWithoutPadding[i]
+          : makeEmptyMembershipWitness(NULLIFIER_TREE_HEIGHT),
+      ),
+      sortedNullifiers: makeTuple(MAX_NULLIFIERS_PER_TX, i => Fr.fromBuffer(sortednullifiers[i])),
+      sortedNullifierIndexes: makeTuple(MAX_NULLIFIERS_PER_TX, i => sortedNewLeavesIndexes[i]),
+      noteHashSubtreeSiblingPath,
+      nullifierSubtreeSiblingPath,
+      lowPublicDataWritesPreimages: padArrayEnd(
+        txPublicDataUpdateRequestInfo.lowPublicDataWritesPreimages,
+        PublicDataTreeLeafPreimage.empty(),
+        MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
+      ),
+      lowPublicDataWritesMembershipWitnesses: padArrayEnd(
+        txPublicDataUpdateRequestInfo.lowPublicDataWritesMembershipWitnesses,
+        MembershipWitness.empty(PUBLIC_DATA_TREE_HEIGHT),
+        MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
+      ),
+      publicDataTreeSiblingPaths: padArrayEnd(
+        txPublicDataUpdateRequestInfo.publicDataWritesSiblingPaths,
+        makeTuple(PUBLIC_DATA_TREE_HEIGHT, () => Fr.ZERO),
+        MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
+      ),
+    });
+
+    const blockHash = tx.constants.historicalHeader.hash();
+    const archiveRootMembershipWitness = await getMembershipWitnessFor(
+      blockHash,
+      MerkleTreeId.ARCHIVE,
+      ARCHIVE_HEIGHT,
+      db,
+    );
 
-  const blockHash = tx.constants.historicalHeader.hash();
-  const archiveRootMembershipWitness = await getMembershipWitnessFor(
-    blockHash,
-    MerkleTreeId.ARCHIVE,
-    ARCHIVE_HEIGHT,
-    db,
-  );
+    return PublicBaseRollupHints.from({
+      start,
+      stateDiffHints,
+      feePayerFeeJuiceBalanceReadHint: feePayerFeeJuiceBalanceReadHint,
+      archiveRootMembershipWitness,
+      constants,
+    });
+  } else {
+    if (
+      txPublicDataUpdateRequestInfo.lowPublicDataWritesMembershipWitnesses.length > 1 ||
+      txPublicDataUpdateRequestInfo.lowPublicDataWritesPreimages.length > 1 ||
+      txPublicDataUpdateRequestInfo.publicDataWritesSiblingPaths.length > 1
+    ) {
+      throw new Error(`More than one public data write in a private only tx`);
+    }
+
+    const feeWriteLowLeafPreimage =
+      txPublicDataUpdateRequestInfo.lowPublicDataWritesPreimages[0] || PublicDataTreeLeafPreimage.empty();
+    const feeWriteLowLeafMembershipWitness =
+      txPublicDataUpdateRequestInfo.lowPublicDataWritesMembershipWitnesses[0] ||
+      MembershipWitness.empty<typeof PUBLIC_DATA_TREE_HEIGHT>(PUBLIC_DATA_TREE_HEIGHT);
+    const feeWriteSiblingPath =
+      txPublicDataUpdateRequestInfo.publicDataWritesSiblingPaths[0] ||
+      makeTuple(PUBLIC_DATA_TREE_HEIGHT, () => Fr.ZERO);
+
+    const stateDiffHints = PrivateBaseStateDiffHints.from({
+      nullifierPredecessorPreimages: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
+        i < nullifierWitnessLeaves.length
+          ? (nullifierWitnessLeaves[i].leafPreimage as NullifierLeafPreimage)
+          : NullifierLeafPreimage.empty(),
+      ),
+      nullifierPredecessorMembershipWitnesses: makeTuple(MAX_NULLIFIERS_PER_TX, i =>
+        i < nullifierPredecessorMembershipWitnessesWithoutPadding.length
+          ? nullifierPredecessorMembershipWitnessesWithoutPadding[i]
+          : makeEmptyMembershipWitness(NULLIFIER_TREE_HEIGHT),
+      ),
+      sortedNullifiers: makeTuple(MAX_NULLIFIERS_PER_TX, i => Fr.fromBuffer(sortednullifiers[i])),
+      sortedNullifierIndexes: makeTuple(MAX_NULLIFIERS_PER_TX, i => sortedNewLeavesIndexes[i]),
+      noteHashSubtreeSiblingPath,
+      nullifierSubtreeSiblingPath,
+      feeWriteLowLeafPreimage,
+      feeWriteLowLeafMembershipWitness,
+      feeWriteSiblingPath,
+    });
+
+    const blockHash = tx.constants.historicalHeader.hash();
+    const archiveRootMembershipWitness = await getMembershipWitnessFor(
+      blockHash,
+      MerkleTreeId.ARCHIVE,
+      ARCHIVE_HEIGHT,
+      db,
+    );
 
-  return BaseRollupHints.from({
-    start,
-    stateDiffHints,
-    feePayerFeeJuiceBalanceReadHint: feePayerFeeJuiceBalanceReadHint,
-    sortedPublicDataWrites: txPublicDataUpdateRequestInfo.sortedPublicDataWrites,
-    sortedPublicDataWritesIndexes: txPublicDataUpdateRequestInfo.sortedPublicDataWritesIndexes,
-    lowPublicDataWritesPreimages: txPublicDataUpdateRequestInfo.lowPublicDataWritesPreimages,
-    lowPublicDataWritesMembershipWitnesses: txPublicDataUpdateRequestInfo.lowPublicDataWritesMembershipWitnesses,
-    archiveRootMembershipWitness,
-    constants,
-  });
+    return PrivateBaseRollupHints.from({
+      start,
+      stateDiffHints,
+      feePayerFeeJuiceBalanceReadHint: feePayerFeeJuiceBalanceReadHint,
+      archiveRootMembershipWitness,
+      constants,
+    });
+  }
 }
 
 async function getPublicDataHint(db: MerkleTreeWriteOperations, leafSlot: bigint) {
@@ -416,67 +480,50 @@ export function makeEmptyMembershipWitness<N extends number>(height: N) {
 }
 
 async function processPublicDataUpdateRequests(tx: ProcessedTx, db: MerkleTreeWriteOperations) {
-  const allPublicDataUpdateRequests = padArrayEnd(
-    tx.txEffect.publicDataWrites,
-    PublicDataWrite.empty(),
-    MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX,
-  );
-
-  const allPublicDataWrites = allPublicDataUpdateRequests.map(
+  const allPublicDataWrites = tx.txEffect.publicDataWrites.map(
     ({ leafSlot, value }) => new PublicDataTreeLeaf(leafSlot, value),
   );
-  const { lowLeavesWitnessData, newSubtreeSiblingPath, sortedNewLeaves, sortedNewLeavesIndexes } = await db.batchInsert(
-    MerkleTreeId.PUBLIC_DATA_TREE,
-    allPublicDataWrites.map(x => x.toBuffer()),
-    // TODO(#3675) remove oldValue from update requests
-    PUBLIC_DATA_SUBTREE_HEIGHT,
-  );
-
-  if (lowLeavesWitnessData === undefined) {
-    throw new Error(`Could not craft public data batch insertion proofs`);
-  }
-
-  const sortedPublicDataWrites = makeTuple(MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX, i => {
-    return PublicDataTreeLeaf.fromBuffer(sortedNewLeaves[i]);
-  });
-
-  const sortedPublicDataWritesIndexes = makeTuple(MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX, i => {
-    return sortedNewLeavesIndexes[i];
-  });
 
-  const subtreeSiblingPathAsFields = newSubtreeSiblingPath.toFields();
-  const newPublicDataSubtreeSiblingPath = makeTuple(PUBLIC_DATA_SUBTREE_SIBLING_PATH_LENGTH, i => {
-    return subtreeSiblingPathAsFields[i];
-  });
+  const lowPublicDataWritesPreimages = [];
+  const lowPublicDataWritesMembershipWitnesses = [];
+  const publicDataWritesSiblingPaths = [];
+
+  for (const write of allPublicDataWrites) {
+    if (write.isEmpty()) {
+      throw new Error(`Empty public data write in tx: ${toFriendlyJSON(tx)}`);
+    }
+
+    // TODO(Alvaro) write a specialized function for this? Internally add_or_update_value uses batch insertion anyway
+    const { lowLeavesWitnessData, newSubtreeSiblingPath } = await db.batchInsert(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      [write.toBuffer()],
+      // TODO(#3675) remove oldValue from update requests
+      0,
+    );
 
-  const lowPublicDataWritesMembershipWitnesses: Tuple<
-    MembershipWitness<typeof PUBLIC_DATA_TREE_HEIGHT>,
-    typeof MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX
-  > = makeTuple(MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX, i => {
-    const witness = lowLeavesWitnessData[i];
-    return MembershipWitness.fromBufferArray(
-      witness.index,
-      assertLength(witness.siblingPath.toBufferArray(), PUBLIC_DATA_TREE_HEIGHT),
+    if (lowLeavesWitnessData === undefined) {
+      throw new Error(`Could not craft public data batch insertion proofs`);
+    }
+
+    const [lowLeafWitness] = lowLeavesWitnessData;
+    lowPublicDataWritesPreimages.push(lowLeafWitness.leafPreimage as PublicDataTreeLeafPreimage);
+    lowPublicDataWritesMembershipWitnesses.push(
+      MembershipWitness.fromBufferArray<typeof PUBLIC_DATA_TREE_HEIGHT>(
+        lowLeafWitness.index,
+        assertLength(lowLeafWitness.siblingPath.toBufferArray(), PUBLIC_DATA_TREE_HEIGHT),
+      ),
     );
-  });
 
-  const lowPublicDataWritesPreimages: Tuple<
-    PublicDataTreeLeafPreimage,
-    typeof MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX
-  > = makeTuple(MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX, i => {
-    return lowLeavesWitnessData[i].leafPreimage as PublicDataTreeLeafPreimage;
-  });
+    const insertionSiblingPath = newSubtreeSiblingPath.toFields();
+    assertLength(insertionSiblingPath, PUBLIC_DATA_TREE_HEIGHT);
 
-  // validate that the sortedPublicDataWrites and sortedPublicDataWritesIndexes are in the correct order
-  // otherwise it will just fail in the circuit
-  assertPermutation(allPublicDataWrites, sortedPublicDataWrites, sortedPublicDataWritesIndexes, (a, b) => a.equals(b));
+    publicDataWritesSiblingPaths.push(insertionSiblingPath as Tuple<Fr, typeof PUBLIC_DATA_TREE_HEIGHT>);
+  }
 
   return {
     lowPublicDataWritesPreimages,
     lowPublicDataWritesMembershipWitnesses,
-    newPublicDataSubtreeSiblingPath,
-    sortedPublicDataWrites,
-    sortedPublicDataWritesIndexes,
+    publicDataWritesSiblingPaths,
   };
 }
 

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

@@ -13,8 +13,10 @@ import {
   AvmProofData,
   type BaseRollupHints,
   Fr,
+  PrivateBaseRollupHints,
   PrivateBaseRollupInputs,
   PrivateTubeData,
+  PublicBaseRollupHints,
   PublicBaseRollupInputs,
   PublicTubeData,
   type TUBE_PROOF_LENGTH,
@@ -66,6 +68,9 @@ export class TxProvingState {
     const vkData = this.getTubeVkData();
     const tubeData = new PrivateTubeData(this.processedTx.data.toKernelCircuitPublicInputs(), this.tube.proof, vkData);
 
+    if (!(this.baseRollupHints instanceof PrivateBaseRollupHints)) {
+      throw new Error('Mismatched base rollup hints, expected private base rollup hints');
+    }
     return new PrivateBaseRollupInputs(tubeData, this.baseRollupHints);
   }
 
@@ -92,6 +97,10 @@ export class TxProvingState {
       this.getAvmVkData(),
     );
 
+    if (!(this.baseRollupHints instanceof PublicBaseRollupHints)) {
+      throw new Error('Mismatched base rollup hints, expected public base rollup hints');
+    }
+
     return new PublicBaseRollupInputs(tubeData, avmProofData, this.baseRollupHints);
   }
 

package/src/proving_broker/proving_agent.ts

@@ -0,0 +1,90 @@
+import {
+  ProvingError,
+  type ProvingRequestType,
+  type ServerCircuitProver,
+  type V2ProvingJob,
+} from '@aztec/circuit-types';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { RunningPromise } from '@aztec/foundation/running-promise';
+
+import { type ProvingJobConsumer } from './proving_broker_interface.js';
+import { ProvingJobController, ProvingJobStatus } from './proving_job_controller.js';
+
+/**
+ * A helper class that encapsulates a circuit prover and connects it to a job source.
+ */
+export class ProvingAgent {
+  private currentJobController?: ProvingJobController;
+  private runningPromise: RunningPromise;
+
+  constructor(
+    /** The source of proving jobs */
+    private jobSource: ProvingJobConsumer,
+    /** The prover implementation to defer jobs to */
+    private circuitProver: ServerCircuitProver,
+    /** Optional list of allowed proof types to build */
+    private proofAllowList?: Array<ProvingRequestType>,
+    /** How long to wait between jobs */
+    private pollIntervalMs = 1000,
+    private log = createDebugLogger('aztec:proving-broker:proving-agent'),
+  ) {
+    this.runningPromise = new RunningPromise(this.safeWork, this.pollIntervalMs);
+  }
+
+  public setCircuitProver(circuitProver: ServerCircuitProver): void {
+    this.circuitProver = circuitProver;
+  }
+
+  public isRunning(): boolean {
+    return this.runningPromise?.isRunning() ?? false;
+  }
+
+  public start(): void {
+    this.runningPromise.start();
+  }
+
+  public async stop(): Promise<void> {
+    this.currentJobController?.abort();
+    await this.runningPromise.stop();
+  }
+
+  private safeWork = async () => {
+    try {
+      // every tick we need to
+      // (1) either do a heartbeat, telling the broker that we're working
+      // (2) get a new job
+      // If during (1) the broker returns a new job that means we can cancel the current job and start the new one
+      let maybeJob: { job: V2ProvingJob; time: number } | undefined;
+      if (this.currentJobController?.getStatus() === ProvingJobStatus.PROVING) {
+        maybeJob = await this.jobSource.reportProvingJobProgress(
+          this.currentJobController.getJobId(),
+          this.currentJobController.getStartedAt(),
+          { allowList: this.proofAllowList },
+        );
+      } else {
+        maybeJob = await this.jobSource.getProvingJob({ allowList: this.proofAllowList });
+      }
+
+      if (!maybeJob) {
+        return;
+      }
+
+      if (this.currentJobController?.getStatus() === ProvingJobStatus.PROVING) {
+        this.currentJobController?.abort();
+      }
+
+      const { job, time } = maybeJob;
+      this.currentJobController = new ProvingJobController(job, time, this.circuitProver, (err, result) => {
+        if (err) {
+          const retry = err.name === ProvingError.NAME ? (err as ProvingError).retry : false;
+          return this.jobSource.reportProvingJobError(job.id, err, retry);
+        } else if (result) {
+          return this.jobSource.reportProvingJobSuccess(job.id, result);
+        }
+      });
+      this.currentJobController.start();
+    } catch (err) {
+      this.log.error(`Error in ProvingAgent: ${String(err)}`);
+    }
+  };
+}