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

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -1,3 +1,4 @@
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { BatchQueue } from '@aztec/foundation/queue';
@@ -10,7 +11,14 @@ import {
   ProvingJobSettledResult,
   getEpochFromProvingJobId,
 } from '@aztec/stdlib/interfaces/server';
-import { Attributes, LmdbMetrics, type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
+import {
+  Attributes,
+  LmdbMetrics,
+  type TelemetryClient,
+  type Tracer,
+  getTelemetryClient,
+  trackSpan,
+} from '@aztec/telemetry-client';
 
 import { mkdir, readdir } from 'fs/promises';
 import { join } from 'path';
@@ -46,9 +54,9 @@ class SingleEpochDatabase {
 
   async *allProvingJobs(): AsyncIterableIterator<[ProvingJob, ProvingJobSettledResult | undefined]> {
     for await (const jobStr of this.jobs.valuesAsync()) {
-      const job = await jsonParseWithSchema(jobStr, ProvingJob);
+      const job = jsonParseWithSchema(jobStr, ProvingJob);
       const resultStr = await this.jobResults.getAsync(job.id);
-      const result = resultStr ? await jsonParseWithSchema(resultStr, ProvingJobSettledResult) : undefined;
+      const result = resultStr ? jsonParseWithSchema(resultStr, ProvingJobSettledResult) : undefined;
       yield [job, result];
     }
   }
@@ -77,6 +85,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
 
   private batchQueue: BatchQueue<ProvingJob | [ProvingJobId, ProvingJobSettledResult], number>;
 
+  public readonly tracer: Tracer;
+
   private constructor(
     private epochs: Map<number, SingleEpochDatabase>,
     private config: ProverBrokerConfig,
@@ -91,6 +101,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       () => this.estimateSize(),
     );
 
+    this.tracer = client.getTracer('KVBrokerDatabase');
+
     this.batchQueue = new BatchQueue(
       (items, key) => this.commitWrites(items, key),
       config.proverBrokerBatchSize,
@@ -104,14 +116,15 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     const jobsToAdd = items.filter((item): item is ProvingJob => 'id' in item);
     const resultsToAdd = items.filter((item): item is [ProvingJobId, ProvingJobSettledResult] => Array.isArray(item));
 
-    const db = await this.getEpochDatabase(epochNumber);
+    const db = await this.getEpochDatabase(EpochNumber(epochNumber));
     await db.batchWrite(jobsToAdd, resultsToAdd);
   }
 
   private async estimateSize() {
     const sizes = await Promise.all(Array.from(this.epochs.values()).map(x => x.estimateSize()));
     return {
-      mappingSize: this.config.dataStoreMapSizeKB,
+      mappingSize: this.config.dataStoreMapSizeKb,
+      physicalFileSize: sizes.reduce((prev, curr) => prev + curr.physicalFileSize, 0),
       numItems: sizes.reduce((prev, curr) => prev + curr.numItems, 0),
       actualSize: sizes.reduce((prev, curr) => prev + curr.actualSize, 0),
     };
@@ -136,13 +149,13 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
         continue;
       }
       logger.info(
-        `Loading broker database for epoch ${epochNumber} from ${fullDirectory} with map size ${config.dataStoreMapSizeKB}KB`,
+        `Loading broker database for epoch ${epochNumber} from ${fullDirectory} with map size ${config.dataStoreMapSizeKb}KB`,
       );
       const db = await openVersionedStoreAt(
         fullDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
         config.l1Contracts.rollupAddress,
-        config.dataStoreMapSizeKB,
+        config.dataStoreMapSizeKb,
       );
       const epochDb = new SingleEpochDatabase(db);
       epochs.set(epochNumber, epochDb);
@@ -163,14 +176,17 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     }
   }
 
-  async deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void> {
-    const oldEpochs = Array.from(this.epochs.keys()).filter(e => e < epochNumber);
+  @trackSpan('KVBrokerDatabase.deleteAllProvingJobsOlderThanEpoch', epochNumber => ({
+    [Attributes.EPOCH_NUMBER]: epochNumber,
+  }))
+  async deleteAllProvingJobsOlderThanEpoch(epochNumber: EpochNumber): Promise<void> {
+    const oldEpochs = Array.from(this.epochs.keys()).filter(e => e < Number(epochNumber));
     for (const old of oldEpochs) {
       const db = this.epochs.get(old);
       if (!db) {
         continue;
       }
-      this.logger.info(`Deleting broker database for epoch ${old}`);
+      this.logger.verbose(`Deleting broker database for epoch ${old}`);
       await db.delete();
       this.epochs.delete(old);
     }
@@ -195,19 +211,19 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     return this.batchQueue.put([id, { status: 'fulfilled', value }], getEpochFromProvingJobId(id));
   }
 
-  private async getEpochDatabase(epochNumber: number): Promise<SingleEpochDatabase> {
+  private async getEpochDatabase(epochNumber: EpochNumber): Promise<SingleEpochDatabase> {
     let epochDb = this.epochs.get(epochNumber);
     if (!epochDb) {
       const newEpochDirectory = join(this.config.dataDirectory!, epochNumber.toString());
       await mkdir(newEpochDirectory, { recursive: true });
       this.logger.info(
-        `Creating broker database for epoch ${epochNumber} at ${newEpochDirectory} with map size ${this.config.dataStoreMapSizeKB}`,
+        `Creating broker database for epoch ${epochNumber} at ${newEpochDirectory} with map size ${this.config.dataStoreMapSizeKb}`,
       );
       const db = await openVersionedStoreAt(
         newEpochDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
         this.config.l1Contracts.rollupAddress,
-        this.config.dataStoreMapSizeKB,
+        this.config.dataStoreMapSizeKb,
       );
       epochDb = new SingleEpochDatabase(db);
       this.epochs.set(epochNumber, epochDb);

package/src/proving_broker/proving_broker_database.ts

@@ -1,3 +1,4 @@
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import type { ProofUri, ProvingJob, ProvingJobId, ProvingJobSettledResult } from '@aztec/stdlib/interfaces/server';
 
 /**
@@ -14,7 +15,7 @@ export interface ProvingBrokerDatabase {
    * Deletes all proving jobs belonging to epochs older than the given epoch
    * @param epochNumber - The epoch number beyond which jobs should be deleted
    */
-  deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void>;
+  deleteAllProvingJobsOlderThanEpoch(epochNumber: EpochNumber): Promise<void>;
 
   /**
    * Returns an iterator over all saved proving jobs

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -8,7 +8,7 @@ import {
   type ObservableResult,
   type TelemetryClient,
   type UpDownCounter,
-  ValueType,
+  createUpDownCounterWithDefault,
 } from '@aztec/telemetry-client';
 
 export type MonitorCallback = (proofType: ProvingRequestType) => number;
@@ -28,49 +28,28 @@ export class ProvingBrokerInstrumentation {
   constructor(client: TelemetryClient, name = 'ProvingBroker') {
     const meter = client.getMeter(name);
 
-    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE, {
-      valueType: ValueType.INT,
-    });
+    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE);
 
-    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS);
 
-    this.resolvedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RESOLVED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    const provingJobTypes = Object.values(ProvingRequestType).filter(v => typeof v === 'string');
+    const provingJobAttrs = { [Attributes.PROVING_JOB_TYPE]: provingJobTypes };
 
-    this.rejectedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_REJECTED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.resolvedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RESOLVED_JOBS, provingJobAttrs);
 
-    this.retriedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RETRIED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.rejectedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_REJECTED_JOBS, provingJobAttrs);
 
-    this.timedOutJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.retriedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RETRIED_JOBS, provingJobAttrs);
 
-    this.cachedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_CACHED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.timedOutJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, provingJobAttrs);
 
-    this.totalJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TOTAL_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.cachedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_CACHED_JOBS, provingJobAttrs);
 
-    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT, {
-      description: 'Records how long a job sits in the queue',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.totalJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TOTAL_JOBS, provingJobAttrs);
 
-    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION, {
-      description: 'Records how long a job takes to complete',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT);
+
+    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION);
   }
 
   monitorQueueDepth(fn: MonitorCallback) {

package/src/proving_broker/proving_job_controller.ts

@@ -1,4 +1,7 @@
-import { createLogger } from '@aztec/foundation/log';
+import { EpochNumber } from '@aztec/foundation/branded-types';
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { AbortError } from '@aztec/foundation/error';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import type {
   ProvingJobId,
   ProvingJobInputs,
@@ -9,38 +12,35 @@ import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
 export enum ProvingJobControllerStatus {
   IDLE = 'idle',
-  PROVING = 'proving',
+  RUNNING = 'running',
   DONE = 'done',
-  ABORTED = 'aborted',
-}
-
-interface ProvingJobCompletionCallback<T extends ProvingRequestType = ProvingRequestType> {
-  (
-    jobId: ProvingJobId,
-    type: T,
-    error: Error | undefined,
-    result: ProvingJobResultsMap[T] | undefined,
-  ): void | Promise<void>;
 }
 
 export class ProvingJobController {
   private status: ProvingJobControllerStatus = ProvingJobControllerStatus.IDLE;
   private promise?: Promise<void>;
   private abortController = new AbortController();
+  private result?: ProvingJobResultsMap[ProvingRequestType] | Error;
+  private log: Logger;
 
   constructor(
     private jobId: ProvingJobId,
     private inputs: ProvingJobInputs,
-    private epochNumber: number,
+    private epochNumber: EpochNumber,
     private startedAt: number,
     private circuitProver: ServerCircuitProver,
-    private onComplete: ProvingJobCompletionCallback,
-    private log = createLogger('prover-client:proving-agent:job-controller'),
-  ) {}
+    private onComplete: () => void,
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('prover-client:proving-agent:job-controller', {
+      instanceId: randomBytes(4).toString('hex'),
+      ...bindings,
+    });
+  }
 
   public start(): void {
     if (this.status !== ProvingJobControllerStatus.IDLE) {
-      this.log.verbose(
+      this.log.warn(
         `Job controller for jobId=${this.jobId} not starting because it is not idle currentStatus=${this.status}`,
         {
           currentStatus: this.status,
@@ -50,63 +50,23 @@ export class ProvingJobController {
       return;
     }
 
-    this.status = ProvingJobControllerStatus.PROVING;
-    this.log.verbose(`Job controller started jobId=${this.jobId}`, {
+    this.promise = this.run();
+
+    this.log.info(`Job controller started jobId=${this.jobId}`, {
       jobId: this.jobId,
     });
-
-    this.promise = this.generateProof()
-      .then(
-        result => {
-          if (this.status === ProvingJobControllerStatus.ABORTED) {
-            this.log.warn(`Job controller for jobId=${this.jobId} completed successfully but job was aborted`, {
-              currentStatus: this.status,
-              jobId: this.jobId,
-            });
-            return;
-          }
-          this.status = ProvingJobControllerStatus.DONE;
-          this.log.verbose(`Job controller for jobId=${this.jobId} completed successfully`, {
-            jobId: this.jobId,
-          });
-          return this.onComplete(this.jobId, this.inputs.type, undefined, result);
-        },
-        error => {
-          if (this.status === ProvingJobControllerStatus.ABORTED) {
-            this.log.warn(`Job controller for jobId=${this.jobId} finished with an error but job was aborted`, {
-              currentStatus: this.status,
-              jobId: this.jobId,
-            });
-            return;
-          }
-
-          if (error.name === 'AbortError') {
-            // Ignore abort errors
-            return;
-          }
-
-          this.log.verbose(`Job controller for jobId=${this.jobId} finished with an error`, {
-            jobId: this.jobId,
-            err: error,
-          });
-
-          this.status = ProvingJobControllerStatus.DONE;
-          return this.onComplete(this.jobId, this.inputs.type, error, undefined);
-        },
-      )
-      .catch(err => {
-        this.log.error(`Job constroller failed to send result for jobId=${this.jobId}: ${err}`, err, {
-          jobId: this.jobId,
-        });
-      });
   }
 
   public getStatus(): ProvingJobControllerStatus {
     return this.status;
   }
 
+  public getResult(): ProvingJobResultsMap[ProvingRequestType] | Error | undefined {
+    return this.result;
+  }
+
   public abort(): void {
-    if (this.status !== ProvingJobControllerStatus.PROVING) {
+    if (this.status !== ProvingJobControllerStatus.RUNNING) {
       this.log.warn(`Tried to abort job controller for jobId=${this.jobId} but it is not running`, {
         currentStatus: this.status,
         jobId: this.jobId,
@@ -114,9 +74,8 @@ export class ProvingJobController {
       return;
     }
 
-    this.status = ProvingJobControllerStatus.ABORTED;
     this.abortController.abort();
-    this.log.verbose(`Aborted job controller for jobId=${this.jobId}`, {
+    this.log.warn(`Aborted job controller for jobId=${this.jobId}`, {
       jobId: this.jobId,
     });
   }
@@ -125,6 +84,10 @@ export class ProvingJobController {
     return this.jobId;
   }
 
+  public getProofType(): ProvingRequestType {
+    return this.inputs.type;
+  }
+
   public getStartedAt(): number {
     return this.startedAt;
   }
@@ -133,6 +96,36 @@ export class ProvingJobController {
     return ProvingRequestType[this.inputs.type];
   }
 
+  private run = async () => {
+    this.status = ProvingJobControllerStatus.RUNNING;
+    let result: ProvingJobResultsMap[ProvingRequestType] | Error;
+    try {
+      result = await this.generateProof();
+    } catch (err) {
+      if (err && err instanceof Error) {
+        result = err;
+      } else {
+        result = new Error('Unknown proving error: ' + String(err), { cause: err });
+      }
+    }
+
+    if (this.abortController.signal.aborted) {
+      this.log.warn(`Job controller for jobId=${this.jobId} completed but job was aborted`, {
+        currentStatus: this.status,
+        jobId: this.jobId,
+      });
+      result = new AbortError('Proof was aborted');
+    }
+
+    this.result = result;
+    this.status = ProvingJobControllerStatus.DONE;
+    try {
+      this.onComplete();
+    } catch (err) {
+      this.log.warn(`On complete handler error: ${err}`, { jobId: this.jobId });
+    }
+  };
+
   private async generateProof(): Promise<ProvingJobResultsMap[ProvingRequestType]> {
     const { type, inputs } = this.inputs;
     const signal = this.abortController.signal;
@@ -141,50 +134,74 @@ export class ProvingJobController {
         return await this.circuitProver.getAvmProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.PRIVATE_BASE_ROLLUP: {
-        return await this.circuitProver.getPrivateBaseRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.PUBLIC_CHONK_VERIFIER: {
+        return await this.circuitProver.getPublicChonkVerifierProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.PRIVATE_TX_BASE_ROLLUP: {
+        return await this.circuitProver.getPrivateTxBaseRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.PUBLIC_TX_BASE_ROLLUP: {
+        return await this.circuitProver.getPublicTxBaseRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.TX_MERGE_ROLLUP: {
+        return await this.circuitProver.getTxMergeRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.PUBLIC_BASE_ROLLUP: {
-        return await this.circuitProver.getPublicBaseRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.MERGE_ROLLUP: {
-        return await this.circuitProver.getMergeRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootSingleTxFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP: {
-        return await this.circuitProver.getEmptyBlockRootRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootEmptyTxFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
       case ProvingRequestType.BLOCK_ROOT_ROLLUP: {
         return await this.circuitProver.getBlockRootRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP: {
-        return await this.circuitProver.getSingleTxBlockRootRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP: {
+        return await this.circuitProver.getBlockRootSingleTxRollupProof(inputs, signal, this.epochNumber);
       }
 
       case ProvingRequestType.BLOCK_MERGE_ROLLUP: {
         return await this.circuitProver.getBlockMergeRollupProof(inputs, signal, this.epochNumber);
       }
 
+      case ProvingRequestType.CHECKPOINT_ROOT_ROLLUP: {
+        return await this.circuitProver.getCheckpointRootRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP: {
+        return await this.circuitProver.getCheckpointRootSingleBlockRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_PADDING_ROLLUP: {
+        return await this.circuitProver.getCheckpointPaddingRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_MERGE_ROLLUP: {
+        return await this.circuitProver.getCheckpointMergeRollupProof(inputs, signal, this.epochNumber);
+      }
+
       case ProvingRequestType.ROOT_ROLLUP: {
         return await this.circuitProver.getRootRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.BASE_PARITY: {
+      case ProvingRequestType.PARITY_BASE: {
         return await this.circuitProver.getBaseParityProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.ROOT_PARITY: {
+      case ProvingRequestType.PARITY_ROOT: {
         return await this.circuitProver.getRootParityProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.TUBE_PROOF: {
-        return await this.circuitProver.getTubeProof(inputs, signal, this.epochNumber);
-      }
-
       default: {
         const _exhaustive: never = type;
         return Promise.reject(new Error(`Invalid proof request type: ${type}`));

package/src/proving_broker/rpc.ts

@@ -1,10 +1,11 @@
+import { EpochNumberSchema } from '@aztec/foundation/branded-types';
 import { createSafeJsonRpcClient } from '@aztec/foundation/json-rpc/client';
-import type { SafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 import {
   type GetProvingJobResponse,
   ProofUri,
   ProvingJob,
   type ProvingJobBroker,
+  type ProvingJobBrokerDebug,
   type ProvingJobConsumer,
   ProvingJobId,
   type ProvingJobProducer,
@@ -13,7 +14,7 @@ import {
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 import { type ApiSchemaFor, optional } from '@aztec/stdlib/schemas';
 import { type ComponentsVersions, getVersioningResponseHandler } from '@aztec/stdlib/versioning';
-import { createTracedJsonRpcServer, makeTracedFetch } from '@aztec/telemetry-client';
+import { makeTracedFetch } from '@aztec/telemetry-client';
 
 import { z } from 'zod';
 
@@ -54,9 +55,17 @@ export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
   ...ProvingJobProducerSchema,
 };
 
-export function createProvingBrokerServer(broker: ProvingJobBroker): SafeJsonRpcServer {
-  return createTracedJsonRpcServer(broker, ProvingJobBrokerSchema);
-}
+export const ProvingJobBrokerDebugSchema: ApiSchemaFor<ProvingJobBrokerDebug> = {
+  replayProvingJob: z
+    .function()
+    .args(ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri)
+    .returns(ProvingJobStatus),
+};
+
+export const ProvingJobBrokerSchemaWithDebug: ApiSchemaFor<ProvingJobBroker & ProvingJobBrokerDebug> = {
+  ...ProvingJobBrokerSchema,
+  ...ProvingJobBrokerDebugSchema,
+};
 
 export function createProvingJobBrokerClient(
   url: string,

package/src/test/mock_proof_store.ts

@@ -0,0 +1,14 @@
+import type { ProvingJobId } from '@aztec/stdlib/interfaces/server';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+
+// Mock ProofStore for faster benchmarks with realistic Cloud Storage URIs
+export class MockProofStore {
+  private mockCounter = 0;
+  private readonly bucketName = 'aztec-proving-benchmarks';
+  private readonly basePath = 'proving-jobs';
+
+  saveProofInput(jobId: ProvingJobId, type: ProvingRequestType): Promise<string> {
+    const uri = `gs://${this.bucketName}/${this.basePath}/inputs/${type}/${jobId}`;
+    return Promise.resolve(uri as any);
+  }
+}