@aztec/prover-client 0.0.0-test.0 → 0.0.1-commit.0208eb9

package/src/proving_broker/proof_store/factory.ts

@@ -1,42 +1,20 @@
 import { createLogger } from '@aztec/foundation/log';
+import { createFileStore } from '@aztec/stdlib/file-store';
 
-import { GoogleCloudStorageProofStore } from './gcs_proof_store.js';
+import { FileStoreProofStore } from './file_store_proof_store.js';
 import { InlineProofStore } from './inline_proof_store.js';
 import type { ProofStore } from './proof_store.js';
 
-export function createProofStore(config: string | undefined, logger = createLogger('prover-client:proof-store')) {
-  if (config === undefined) {
+export async function createProofStore(
+  config: string | undefined,
+  logger = createLogger('prover-client:proof-store'),
+): Promise<ProofStore> {
+  if (!config) {
     logger.info('Creating inline proof store');
     return new InlineProofStore();
-  } else if (config.startsWith('gs://')) {
-    try {
-      const url = new URL(config);
-      const bucket = url.host;
-      const path = url.pathname.replace(/^\/+/, '');
-      logger.info(`Creating google cloud proof store at ${bucket}`, { bucket, path });
-      return new GoogleCloudStorageProofStore(bucket, path);
-    } catch (err) {
-      throw new Error(
-        `Invalid google cloud proof store definition: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`,
-      );
-    }
-  } else {
-    throw new Error(`Unknown proof store config: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`);
   }
-}
 
-export function createProofStoreForUri(
-  uri: string,
-  logger = createLogger('prover-client:proof-store'),
-): Pick<ProofStore, 'getProofInput' | 'getProofOutput'> {
-  if (uri.startsWith('data://')) {
-    return createProofStore(undefined, logger);
-  } else if (uri.startsWith('gs://')) {
-    const url = new URL(uri);
-    const basePath = url.pathname.replace(/^\/+/, '').split('/').slice(0, -3);
-    url.pathname = basePath.join('/');
-    return createProofStore(uri, logger);
-  } else {
-    throw new Error(`Unknown proof store config: '${uri}'. Supported protocols are 'data://' and 'gs://'.`);
-  }
+  const fileStore = await createFileStore(config, logger);
+  logger.info(`Creating file store proof store at ${config}`);
+  return new FileStoreProofStore(fileStore);
 }

package/src/proving_broker/proof_store/file_store_proof_store.ts

@@ -0,0 +1,78 @@
+import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
+import type { FileStore } from '@aztec/stdlib/file-store';
+import {
+  type ProofUri,
+  type ProvingJobId,
+  type ProvingJobInputs,
+  type ProvingJobInputsMap,
+  ProvingJobResult,
+  type ProvingJobResultsMap,
+  getProvingJobInputClassFor,
+} from '@aztec/stdlib/interfaces/server';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+
+import type { ProofStore } from './proof_store.js';
+
+const INPUTS_PATH = 'inputs';
+const OUTPUTS_PATH = 'outputs';
+
+/**
+ * A proof store implementation backed by a generic FileStore.
+ * Supports any storage backend (GCS, S3, local filesystem) via the FileStore abstraction.
+ */
+export class FileStoreProofStore implements ProofStore {
+  constructor(private readonly fileStore: FileStore) {}
+
+  async saveProofInput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri> {
+    const path = `${INPUTS_PATH}/${ProvingRequestType[type]}/${id}`;
+    const uri = await this.fileStore.save(path, inputs.toBuffer());
+    return uri as ProofUri;
+  }
+
+  async saveProofOutput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri> {
+    const jobResult = { type, result } as ProvingJobResult;
+    const json = jsonStringify(jobResult);
+    const path = `${OUTPUTS_PATH}/${ProvingRequestType[type]}/${id}.json`;
+    const uri = await this.fileStore.save(path, Buffer.from(json, 'utf-8'));
+    return uri as ProofUri;
+  }
+
+  async getProofInput(uri: ProofUri): Promise<ProvingJobInputs> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      const type = this.extractTypeFromUri(uri);
+      const inputs = getProvingJobInputClassFor(type).fromBuffer(buffer);
+      return { inputs, type } as ProvingJobInputs;
+    } catch (err) {
+      throw new Error(`Error getting proof input at ${uri}: ${err}`);
+    }
+  }
+
+  async getProofOutput(uri: ProofUri): Promise<ProvingJobResult> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      return jsonParseWithSchema(buffer.toString('utf-8'), ProvingJobResult);
+    } catch (err) {
+      throw new Error(`Error getting proof output at ${uri}: ${err}`);
+    }
+  }
+
+  private extractTypeFromUri(uri: string): ProvingRequestType {
+    const url = new URL(uri);
+    const pathParts = url.pathname.split('/').filter(Boolean);
+    const typeString = pathParts.at(-2);
+    const type = typeString ? ProvingRequestType[typeString as keyof typeof ProvingRequestType] : undefined;
+    if (type === undefined) {
+      throw new Error(`Unrecognized proof type ${typeString} in URI ${uri}`);
+    }
+    return type;
+  }
+}

package/src/proving_broker/proof_store/index.ts

@@ -1,3 +1,4 @@
 export * from './proof_store.js';
 export * from './inline_proof_store.js';
 export * from './factory.js';
+export * from './file_store_proof_store.js';

package/src/proving_broker/proof_store/inline_proof_store.ts

@@ -52,7 +52,7 @@ export class InlineProofStore implements ProofStore {
     return (PREFIX + SEPARATOR + encoded) as ProofUri;
   }
 
-  private decode<T>(uri: ProofUri, schema: ZodFor<T>): Promise<T> {
+  private decode<T>(uri: ProofUri, schema: ZodFor<T>): T {
     const [prefix, data] = uri.split(SEPARATOR);
     if (prefix !== PREFIX) {
       throw new Error('Invalid proof input URI: ' + prefix);

package/src/proving_broker/proving_agent.ts

@@ -1,10 +1,11 @@
-import { createLogger } from '@aztec/foundation/log';
+import { AbortError } from '@aztec/foundation/error';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { truncate } from '@aztec/foundation/string';
-import { Timer } from '@aztec/foundation/timer';
 import { ProvingError } from '@aztec/stdlib/errors';
 import type {
-  ProvingJob,
+  GetProvingJobResponse,
+  ProverAgentStatus,
   ProvingJobConsumer,
   ProvingJobId,
   ProvingJobInputs,
@@ -12,28 +13,17 @@ import type {
   ServerCircuitProver,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
-import {
-  type TelemetryClient,
-  type Traceable,
-  type Tracer,
-  getTelemetryClient,
-  trackSpan,
-} from '@aztec/telemetry-client';
 
 import type { ProofStore } from './proof_store/index.js';
-import { ProvingAgentInstrumentation } from './proving_agent_instrumentation.js';
 import { ProvingJobController, ProvingJobControllerStatus } from './proving_job_controller.js';
 
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProvingAgent implements Traceable {
+export class ProvingAgent {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
-  private instrumentation: ProvingAgentInstrumentation;
-  private idleTimer: Timer | undefined;
-
-  public readonly tracer: Tracer;
+  private log: Logger;
 
   constructor(
     /** The source of proving jobs */
@@ -46,12 +36,9 @@ export class ProvingAgent implements Traceable {
     private proofAllowList: Array<ProvingRequestType> = [],
     /** How long to wait between jobs */
     private pollIntervalMs = 1000,
-    /** A telemetry client through which to emit metrics */
-    client: TelemetryClient = getTelemetryClient(),
-    private log = createLogger('prover-client:proving-agent'),
+    bindings?: LoggerBindings,
   ) {
-    this.tracer = client.getTracer('ProvingAgent');
-    this.instrumentation = new ProvingAgentInstrumentation(client);
+    this.log = createLogger('prover-client:proving-agent', bindings);
     this.runningPromise = new RunningPromise(this.work.bind(this), this.log, this.pollIntervalMs);
   }
 
@@ -64,7 +51,6 @@ export class ProvingAgent implements Traceable {
   }
 
   public start(): void {
-    this.idleTimer = new Timer();
     this.runningPromise.start();
   }
 
@@ -73,41 +59,77 @@ export class ProvingAgent implements Traceable {
     await this.runningPromise.stop();
   }
 
-  @trackSpan('ProvingAgent.safeWork')
+  public getStatus(): ProverAgentStatus {
+    if (this.currentJobController) {
+      return {
+        status: 'proving',
+        jobId: this.currentJobController.getJobId(),
+        proofType: this.currentJobController.getProofType(),
+        startedAtISO: new Date(this.currentJobController.getStartedAt()).toISOString(),
+      };
+    }
+
+    return this.runningPromise.isRunning() ? { status: 'running' } : { status: 'stopped' };
+  }
+
   private async work() {
-    // 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: ProvingJob; time: number } | undefined;
-    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
-      maybeJob = await this.broker.reportProvingJobProgress(
-        this.currentJobController.getJobId(),
-        this.currentJobController.getStartedAt(),
-        { allowList: this.proofAllowList },
-      );
+    // every tick we need to take one of the following actions:
+    // 1. send a hearbeat to the broker that we're working on some job
+    // 2. if the job is complete, send its result to the broker
+    // 3. get a job from the broker
+    // Any one of these actions could give us a new job to work on. If that happens we abort the current job.
+    //
+    // This loop gets triggered in one of two ways:
+    // - either on a timer (see pollIntervalMs)
+    // - or when a proof completes
+    let maybeJob: GetProvingJobResponse | undefined;
+
+    if (this.currentJobController) {
+      const status = this.currentJobController.getStatus();
+      const jobId = this.currentJobController.getJobId();
+      const proofType = this.currentJobController.getProofType();
+      const startedAt = this.currentJobController.getStartedAt();
+      const result = this.currentJobController.getResult();
+
+      if (status === ProvingJobControllerStatus.RUNNING) {
+        maybeJob = await this.broker.reportProvingJobProgress(jobId, startedAt, { allowList: this.proofAllowList });
+      } else if (status === ProvingJobControllerStatus.DONE) {
+        if (result) {
+          maybeJob = await this.reportResult(jobId, proofType, result);
+        } else {
+          this.log.warn(
+            `Job controller for job ${this.currentJobController.getJobId()} is done but doesn't have a result`,
+            { jobId },
+          );
+          maybeJob = await this.reportResult(
+            jobId,
+            proofType,
+            new ProvingError('No result found after proving', undefined, /* retry */ true),
+          );
+        }
+
+        this.currentJobController = undefined;
+      } else {
+        // IDLE status should not be seen because a job is started as soon as it is created
+        this.log.warn(`Idle job controller for job: ${this.currentJobController.getJobId()}. Skipping main loop work`, {
+          jobId: this.currentJobController.getJobId(),
+        });
+        return;
+      }
     } else {
       maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
     }
 
-    if (!maybeJob) {
-      return;
-    }
-
-    if (this.idleTimer) {
-      this.instrumentation.recordIdleTime(this.idleTimer);
+    if (maybeJob) {
+      await this.startJob(maybeJob);
     }
-    this.idleTimer = undefined;
-
-    const { job, time } = maybeJob;
-    await this.startJob(job, time);
   }
 
-  private async startJob(job: ProvingJob, startedAt: number): Promise<void> {
+  private async startJob({ job, time: startedAt }: GetProvingJobResponse): Promise<void> {
     let abortedProofJobId: string | undefined;
     let abortedProofName: string | undefined;
 
-    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
+    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.RUNNING) {
       abortedProofJobId = this.currentJobController.getJobId();
       abortedProofName = this.currentJobController.getProofTypeName();
       this.currentJobController?.abort();
@@ -116,13 +138,13 @@ export class ProvingAgent implements Traceable {
     let inputs: ProvingJobInputs;
     try {
       inputs = await this.proofStore.getProofInput(job.inputsUri);
-    } catch (err) {
+    } catch {
       const maybeJob = await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true, {
         allowList: this.proofAllowList,
       });
 
       if (maybeJob) {
-        return this.startJob(maybeJob.job, maybeJob.time);
+        return this.startJob(maybeJob);
       }
 
       return;
@@ -134,7 +156,12 @@ export class ProvingAgent implements Traceable {
       job.epochNumber,
       startedAt,
       this.circuitProver,
-      this.handleJobResult,
+      () => {
+        // trigger a run of the main work loop when proving completes
+        // no need to await this here. The controller will stay alive (in DONE state) until the result is send to the broker
+        void this.runningPromise.trigger();
+      },
+      this.log.getBindings(),
     );
 
     if (abortedProofJobId) {
@@ -154,28 +181,30 @@ export class ProvingAgent implements Traceable {
     this.currentJobController.start();
   }
 
-  handleJobResult = async <T extends ProvingRequestType>(
+  private async reportResult<T extends ProvingRequestType>(
     jobId: ProvingJobId,
     type: T,
-    err: Error | undefined,
-    result: ProvingJobResultsMap[T] | undefined,
-  ) => {
-    let maybeJob: { job: ProvingJob; time: number } | undefined;
-    if (err) {
-      const retry = err.name === ProvingError.NAME ? (err as ProvingError).retry : false;
-      this.log.error(`Job id=${jobId} type=${ProvingRequestType[type]} failed err=${err.message} retry=${retry}`, err);
-      maybeJob = await this.broker.reportProvingJobError(jobId, err.message, retry, { allowList: this.proofAllowList });
-    } else if (result) {
+    result: ProvingJobResultsMap[T] | Error,
+  ): Promise<GetProvingJobResponse | undefined> {
+    let maybeJob: GetProvingJobResponse | undefined;
+    if (result instanceof AbortError) {
+      // no-op
+      this.log.warn(`Job id=${jobId} was aborted. Not reporting result back to broker`, result);
+    } else if (result instanceof Error) {
+      const retry = result.name === ProvingError.NAME ? (result as ProvingError).retry : false;
+      this.log.error(
+        `Job id=${jobId} type=${ProvingRequestType[type]} failed err=${result.message} retry=${retry}`,
+        result,
+      );
+      maybeJob = await this.broker.reportProvingJobError(jobId, result.message, retry, {
+        allowList: this.proofAllowList,
+      });
+    } else {
       const outputUri = await this.proofStore.saveProofOutput(jobId, type, result);
       this.log.info(`Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncate(outputUri)}`);
       maybeJob = await this.broker.reportProvingJobSuccess(jobId, outputUri, { allowList: this.proofAllowList });
     }
 
-    if (maybeJob) {
-      const { job, time } = maybeJob;
-      await this.startJob(job, time);
-    } else {
-      this.idleTimer = new Timer();
-    }
-  };
+    return maybeJob;
+  }
 }

package/src/proving_broker/proving_broker.ts

@@ -1,26 +1,23 @@
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import { createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { PriorityMemoryQueue } from '@aztec/foundation/queue';
 import { Timer } from '@aztec/foundation/timer';
-import type {
-  GetProvingJobResponse,
-  ProofUri,
-  ProvingJob,
-  ProvingJobConsumer,
-  ProvingJobFilter,
-  ProvingJobId,
-  ProvingJobProducer,
-  ProvingJobSettledResult,
-  ProvingJobStatus,
+import {
+  type GetProvingJobResponse,
+  type ProofUri,
+  type ProvingJob,
+  type ProvingJobBrokerDebug,
+  type ProvingJobConsumer,
+  type ProvingJobFilter,
+  type ProvingJobId,
+  type ProvingJobProducer,
+  type ProvingJobSettledResult,
+  type ProvingJobStatus,
+  tryStop,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
-import {
-  type TelemetryClient,
-  type Traceable,
-  type Tracer,
-  getTelemetryClient,
-  trackSpan,
-} from '@aztec/telemetry-client';
+import { type TelemetryClient, type Traceable, type Tracer, getTelemetryClient } from '@aztec/telemetry-client';
 
 import assert from 'assert';
 
@@ -40,23 +37,36 @@ type EnqueuedProvingJob = Pick<ProvingJob, 'id' | 'epochNumber'>;
  * A broker that manages proof requests and distributes them to workers based on their priority.
  * It takes a backend that is responsible for storing and retrieving proof requests and results.
  */
-export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Traceable {
+export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, ProvingJobBrokerDebug, Traceable {
   private queues: ProvingQueues = {
     [ProvingRequestType.PUBLIC_VM]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.TUBE_PROOF]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PUBLIC_CHONK_VERIFIER]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.PRIVATE_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.PUBLIC_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PRIVATE_TX_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PUBLIC_TX_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.TX_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
     [ProvingRequestType.BLOCK_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
+    [ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
     [ProvingRequestType.BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.BASE_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.ROOT_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.CHECKPOINT_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
+    [ProvingRequestType.CHECKPOINT_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.CHECKPOINT_PADDING_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+
+    [ProvingRequestType.PARITY_BASE]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PARITY_ROOT]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
   };
 
   // holds a copy of the database in memory in order to quickly fulfill requests
@@ -93,7 +103,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   /**
    * The broker keeps track of the highest epoch its seen.
    * This information is used for garbage collection: once it reaches the next epoch, it can start pruning the database of old state.
-   * It is important that this value is initialised to zero. This ensures that we don't delete any old jobs until the current
+   * It is important that this value is initialized to zero. This ensures that we don't delete any old jobs until the current
    * process instance receives a job request informing it of the actual current highest epoch
    * Example:
    * proving epoch 11 - the broker will wipe all jobs for epochs 9 and lower
@@ -105,6 +115,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
   private started = false;
 
+  private debugReplayEnabled: boolean;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     {
@@ -112,6 +124,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       proverBrokerPollIntervalMs,
       proverBrokerJobMaxRetries,
       proverBrokerMaxEpochsToKeepResultsFor,
+      proverBrokerDebugReplayEnabled,
     }: Required<
       Pick<
         ProverBrokerConfig,
@@ -119,6 +132,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         | 'proverBrokerPollIntervalMs'
         | 'proverBrokerJobMaxRetries'
         | 'proverBrokerMaxEpochsToKeepResultsFor'
+        | 'proverBrokerDebugReplayEnabled'
       >
     > = defaultProverBrokerConfig,
     client: TelemetryClient = getTelemetryClient(),
@@ -130,6 +144,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.jobTimeoutMs = proverBrokerJobTimeoutMs!;
     this.maxRetries = proverBrokerJobMaxRetries!;
     this.maxEpochsToKeepResultsFor = proverBrokerMaxEpochsToKeepResultsFor!;
+    this.debugReplayEnabled = proverBrokerDebugReplayEnabled ?? false;
   }
 
   private measureQueueDepth: MonitorCallback = (type: ProvingRequestType) => {
@@ -184,7 +199,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       this.logger.warn('ProvingBroker not started');
       return Promise.resolve();
     }
-    await this.cleanupPromise.stop();
+    await tryStop(this.cleanupPromise);
   }
 
   public enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
@@ -232,6 +247,29 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return Promise.resolve(this.#reportProvingJobProgress(id, startedAt, filter));
   }
 
+  public async replayProvingJob(
+    jobId: ProvingJobId,
+    type: ProvingRequestType,
+    epochNumber: EpochNumber,
+    inputsUri: ProofUri,
+  ): Promise<ProvingJobStatus> {
+    if (!this.debugReplayEnabled) {
+      throw new Error('Debug replay not enabled. Set PROVER_BROKER_DEBUG_REPLAY_ENABLED=true');
+    }
+
+    this.logger.info(`Replaying proving job`, { provingJobId: jobId, epochNumber, inputsUri });
+
+    // Clear existing state and enqueue
+    this.cleanUpProvingJobState([jobId]);
+
+    const job: ProvingJob = { id: jobId, type, epochNumber, inputsUri };
+    this.jobsCache.set(jobId, job);
+    await this.database.addProvingJob(job);
+    this.enqueueJobInternal(job);
+
+    return { status: 'in-queue' };
+  }
+
   async #enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
     // We return the job status at the start of this call
     const jobStatus = this.#getProvingJobStatus(job.id);
@@ -287,6 +325,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       this.resultsCache.delete(id);
       this.inProgress.delete(id);
       this.retries.delete(id);
+      this.enqueuedAt.delete(id);
     }
   }
 
@@ -313,7 +352,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return Promise.resolve(notifications.concat(completedJobs));
   }
 
-  // eslint-disable-next-line require-await
   #getProvingJob(filter: ProvingJobFilter = { allowList: [] }): { job: ProvingJob; time: number } | undefined {
     const allowedProofs: ProvingRequestType[] =
       Array.isArray(filter.allowList) && filter.allowList.length > 0
@@ -340,6 +378,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
           const enqueuedAt = this.enqueuedAt.get(job.id);
           if (enqueuedAt) {
             this.instrumentation.recordJobWait(job.type, enqueuedAt);
+            // we can clear this flag now.
+            this.enqueuedAt.delete(job.id);
           }
 
           return { job, time };
@@ -548,13 +588,12 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return this.#getProvingJob(filter);
   }
 
-  @trackSpan('ProvingBroker.cleanupPass')
   private async cleanupPass() {
     this.cleanupStaleJobs();
     this.reEnqueueExpiredJobs();
     const oldestEpochToKeep = this.oldestEpochToKeep();
     if (oldestEpochToKeep > 0) {
-      await this.database.deleteAllProvingJobsOlderThanEpoch(oldestEpochToKeep);
+      await this.database.deleteAllProvingJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       this.logger.trace(`Deleted all epochs older than ${oldestEpochToKeep}`);
     }
   }
@@ -571,7 +610,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
     if (jobsToClean.length > 0) {
       this.cleanUpProvingJobState(jobsToClean);
-      this.logger.info(`Cleaned up jobs=${jobsToClean.length}`);
+      this.logger.verbose(`Cleaned up proving jobs=${jobsToClean.length}`);
     }
   }
 
@@ -671,17 +710,23 @@ function proofTypeComparator(a: ProvingRequestType, b: ProvingRequestType): -1 |
  * The aim is that this will speed up block proving as the closer we get to a block's root proof the more likely it
  * is to get picked up by agents
  */
-const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
+export const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
+  ProvingRequestType.ROOT_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP,
   ProvingRequestType.BLOCK_ROOT_ROLLUP,
-  ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP,
   ProvingRequestType.BLOCK_MERGE_ROLLUP,
-  ProvingRequestType.ROOT_ROLLUP,
-  ProvingRequestType.MERGE_ROLLUP,
-  ProvingRequestType.PUBLIC_BASE_ROLLUP,
-  ProvingRequestType.PRIVATE_BASE_ROLLUP,
+  ProvingRequestType.CHECKPOINT_ROOT_ROLLUP,
+  ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP,
+  ProvingRequestType.CHECKPOINT_MERGE_ROLLUP,
+  ProvingRequestType.CHECKPOINT_PADDING_ROLLUP,
+  ProvingRequestType.TX_MERGE_ROLLUP,
+  ProvingRequestType.PUBLIC_TX_BASE_ROLLUP,
+  ProvingRequestType.PRIVATE_TX_BASE_ROLLUP,
   ProvingRequestType.PUBLIC_VM,
-  ProvingRequestType.TUBE_PROOF,
-  ProvingRequestType.ROOT_PARITY,
-  ProvingRequestType.BASE_PARITY,
-  ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
+  ProvingRequestType.PUBLIC_CHONK_VERIFIER,
+  ProvingRequestType.PARITY_ROOT,
+  ProvingRequestType.PARITY_BASE,
 ];

package/src/proving_broker/proving_broker_database/memory.ts

@@ -1,3 +1,4 @@
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import {
   type ProofUri,
   type ProvingJob,
@@ -43,7 +44,7 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
     return Promise.resolve();
   }
 
-  deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void> {
+  deleteAllProvingJobsOlderThanEpoch(epochNumber: EpochNumber): Promise<void> {
     const toDelete = [
       ...Array.from(this.jobs.keys()).filter(x => getEpochFromProvingJobId(x) < epochNumber),
       ...Array.from(this.results.keys()).filter(x => getEpochFromProvingJobId(x) < epochNumber),
@@ -52,7 +53,7 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
   }
 
   async *allProvingJobs(): AsyncIterableIterator<[ProvingJob, ProvingJobSettledResult | undefined]> {
-    for await (const item of this.jobs.values()) {
+    for (const item of this.jobs.values()) {
       yield [item, this.results.get(item.id)] as const;
     }
   }